from SmolLm3 import LlamaModel
import torch
import yaml
from transformers import AutoTokenizer
from torch.utils.data import DataLoader
import numpy as np
from datasets import load_dataset
import logging
import math
from utils import upload_file_to_s3
# At the start of training loop
# print(f"GPU Memory allocated: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
# print(f"GPU Memory reserved: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")
logger = logging.getLogger(__name__)
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler('training.log')
file_handler.setFormatter(formatter) # Set formatter on the handler, not the logger
logger.addHandler(file_handler)
logger.setLevel(logging.INFO)
def encode_text(examples, tokenizer, seq_length):
"""Tokenize and prepare text examples for training."""
tokens = tokenizer(
examples["text"],
truncation=True,
padding="max_length",
max_length=seq_length + 1,
return_tensors="pt",
)
# Use clone().detach() as recommended
input_ids = tokens["input_ids"].squeeze(0).clone().detach()
input_ids = torch.clamp(input_ids, min=0, max=tokenizer.vocab_size - 1)
labels = input_ids.clone().detach()
labels = labels[1:].to(torch.int64)
input_ids = input_ids[:-1].to(torch.int64)
return {"input_ids": input_ids, "labels": labels}
def load_cosmopedia_dataset(batch_size=8, seq_length=1024, tokenizer=None):
"""
Returns a torch dataloader for the cosmopedia dataset
"""
# Set tokenizer parallelism explicitly
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
logger.info("tokenizer parallelism set to false")
try:
# Increase timeout and retries for dataset loading
from datasets import config
config.HF_DATASETS_TIMEOUT = 300 # 5 minutes timeout
config.MAX_RETRIES = 10 # Increase retry attempts
logger.info("dataset loading config set")
train_dataset = load_dataset(
"HuggingFaceTB/smollm-corpus",
name="cosmopedia-v2",
split="train",
streaming=True,
)
logger.info("dataset loaded")
# Use partial to bind tokenizer and seq_length to the encode function
from functools import partial
encode_fn = partial(encode_text, tokenizer=tokenizer, seq_length=seq_length)
train_dataset = train_dataset.map(
encode_fn,
remove_columns=["text"],
batched=False
)
train_dataset = train_dataset.with_format("torch")
train_dataloader = DataLoader(
train_dataset,
batch_size=batch_size,
num_workers=2,
pin_memory=True,
prefetch_factor=4,
persistent_workers=True
)
return train_dataloader
except Exception as e:
logger.error(f"Error loading dataset: {str(e)}")
return None
def generate(model, idx, max_new_tokens, context_length, temperature=1.0, top_k=None, eos_token=None, device=None):
logger.info(f"Generating on device {device}")
model = model.to(device)
idx = idx.to(device)
model.eval()
for _ in range(max_new_tokens):
idx_cond = idx[:, -context_length:]
with torch.no_grad():
logits, _ = model(idx_cond) # Unpack both logits and loss (ignore loss)
logits = logits.view(idx_cond.shape[0], -1, model.config['vocab_size']) # Reshape to [batch, seq, vocab]
# Get the logits for the last token only
logits = logits[:, -1, :] # Shape: [batch_size, vocab_size]
if top_k is not None:
# top k sampling
top_logits, top_pos = torch.topk(logits, top_k)
min_logit = top_logits[:, -1].unsqueeze(-1)
logits = torch.where(logits < min_logit,
torch.tensor(float('-inf')).to(logits.device),
logits)
# temperature scaling
if temperature > 0.0:
logits /= temperature
probs = torch.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
else:
idx_next = torch.argmax(logits, dim=-1, keepdim=True)
if idx_next.item() == eos_token:
break
idx = torch.cat((idx, idx_next), dim=1)
model.train()
return idx
def sync_device(device):
if device.startswith('cuda'):
torch.cuda.synchronize()
elif device == 'cpu':
torch.cpu.synchronize() if hasattr(torch.cpu, 'synchronize') else None
elif device.startswith('mps'): # For Apple Silicon
torch.mps.synchronize()
def print_gpu_memory(step_name=""):
"""
Print GPU memory statistics with a specified step name
"""
if torch.cuda.is_available():
logger.info(f"\nGPU Memory Stats {step_name}:")
logger.info(f"GPU Memory allocated: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
logger.info(f"GPU Memory reserved: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")
logger.info(f"Max GPU Memory allocated: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")
# Learning rate scheduler
def get_lr_lambda(current_step, warmup_steps, max_steps, max_lr):
"""
Modified learning rate scheduler with:
1. Linear warmup for first 3000 steps
2. Cosine decay from 3000 to 60000 steps
3. Minimum learning rate of 1.5e-5 (5% of max_lr)
"""
min_lr = max_lr * 0.05 # Minimum learning rate (5% of max_lr)
if current_step < warmup_steps:
# Linear warmup from 0 to max_lr
return float(current_step) / float(max(1, warmup_steps))
else:
# Cosine decay from max_lr to min_lr
progress = float(current_step - warmup_steps) / float(max(1, max_steps - warmup_steps))
return min_lr + 0.5 * (max_lr - min_lr) * (1.0 + math.cos(math.pi * progress))
def train_model(config, model, train_loader, test_loader, optimizer, device, num_epochs, eval_freq, eval_iter, start_context="Jack Gisburn rather a cheap genius- ", tokenizer=None):
total_loss = 0
tokens_seen, global_step = 0, -1
# Adjusted gradient accumulation setup
actual_batch_size = config['tokens']['micro_batch_size'] # Now 16
effective_batch_size_multiplier = 2 # Reduced from 4 to maintain reasonable memory usage
target_batch_size = effective_batch_size_multiplier * config['tokens']['micro_batch_size']
gradient_accumulation_steps = target_batch_size // actual_batch_size
# Adjusted learning rate parameters for new batch size
max_lr = 3e-4 # Keep the same max learning rate
warmup_steps = 3000 # Increase warmup steps for longer training
max_steps = 60000 # Set to match 10 hours of training
min_lr = max_lr * 0.05 # Reduce minimum LR to 5% of max (was 10%)
# Create LambdaLR scheduler with the improved lambda function
lr_lambda = lambda step: get_lr_lambda(step, warmup_steps, max_steps, max_lr)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
logger.info(f"Training with learning rate schedule:")
logger.info(f"Max LR: {max_lr}")
logger.info(f"Warmup Steps: {warmup_steps}")
logger.info(f"Max Steps: {max_steps}")
logger.info(f"Min LR: {max_lr * 0.05}")
logger.info(f"Gradient Accumulation Steps: {gradient_accumulation_steps}")
logger.info(f"Effective Batch Size: {actual_batch_size * gradient_accumulation_steps}")
print_gpu_memory("at start of training")
# Add these near the start of training loop
torch.cuda.empty_cache()
torch.backends.cudnn.benchmark = True
for epoch in range(num_epochs):
model.train()
optimizer.zero_grad() # Zero gradients at start of epoch
for batch_idx, batch in enumerate(train_loader):
input_batch = batch['input_ids'].to(device)
target_batch = batch['labels'].to(device)
# Forward pass
with torch.autocast(device_type=device, dtype=torch.bfloat16):
logits, original_loss = model(input_batch, target_batch)
# Scale loss for gradient accumulation
scaled_loss = original_loss / gradient_accumulation_steps
scaled_loss.backward()
# Add the original loss to total_loss for logging
total_loss += original_loss.item() # Don't multiply back up
tokens_seen += input_batch.numel()
# Calculate running average loss
total_batches = batch_idx + 1
avg_loss = total_loss / total_batches
if batch_idx % 25 == 0:
logger.info(f"Batch {batch_idx + 1}, Running Avg Loss: {avg_loss:.5f}")
# Only update weights after accumulating gradients
if (batch_idx + 1) % gradient_accumulation_steps == 0:
# Gradient clipping
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
scheduler.step() # Update learning rate
optimizer.zero_grad()
global_step += 1
# Evaluation block
if global_step % eval_freq == 0 and global_step > 0:
# Use total batches processed instead of global_step
current_lr = scheduler.get_last_lr()[0]
optimizer_lr = optimizer.param_groups[0]['lr']
print_gpu_memory(f"at step {global_step}")
logger.info(f"learning rate: {current_lr:.8f}")
logger.info(f"Ep {epoch+1} (Step {global_step:06d}): "
f"Avg loss {avg_loss:.3f} | {tokens_seen} tokens seen")
logger.info(f"optimizer lr: {optimizer_lr:.8f}")
logger.info(f"scheduler lr: {current_lr:.8f}")
# Generate sample text
encoded_text = tokenizer.encode(start_context, return_tensors="pt")
random_topk = np.random.randint(1, 10)
logger.info(f"random_topk: {random_topk}")
random_temperature = np.random.uniform(0.7, 0.9)
logger.info(f"random_temperature: {random_temperature}")
logger.info(f"global step {global_step} , batch_idx {batch_idx} => generating text")
generated_text = generate(model,
idx=encoded_text,
max_new_tokens=256,
context_length=256,
temperature=random_temperature,
top_k=random_topk,
eos_token=tokenizer.eos_token_id,
device=device)
logger.info(f"+++"*30)
logger.info(tokenizer.decode(generated_text.squeeze(0)))
logger.info(f"+++"*30)
# Save checkpoint
model_file_name = f"model_{global_step}_steps_avg_loss_{avg_loss:.5f}_optimizer_lr_{optimizer_lr:.8f}.pth"
torch.save({
'step': global_step,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'loss': avg_loss,
}, model_file_name)
s3_path = upload_file_to_s3(model_file_name, config['model']['model_config']['s3_bucket'],
config['model']['model_config']['s3_checkpoint_folder'])
logger.info(f"Model saved to S3: {s3_path}")
log_path = upload_file_to_s3(config['model']['model_config']['s3_log_file_name'], config['model']['model_config']['s3_bucket'],
config['model']['model_config']['s3_log_folder'])
logger.info(f"Log saved to S3: {log_path}")
if batch_idx % 100 == 0:
logger.info(f"Batch {batch_idx} finished")
logger.info(f"+++"*30)
logger.info("Training complete")
if __name__ == "__main__":
config = yaml.load(open("config_smollm2_135M.yaml", "r"), Loader=yaml.FullLoader)
logger.info(config)
# Set memory efficient settings
torch.set_float32_matmul_precision('high')
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
# Empty cache before model creation
torch.cuda.empty_cache()
model = LlamaModel(config['model'])
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# Enable gradient checkpointing for memory efficiency
# model.gradient_checkpointing_enable()
model.to(device)
model = torch.compile(model)
logger.info(model)
logger.info("++"*30)
optimizer = torch.optim.AdamW(
model.parameters(),
lr=3e-4,
weight_decay=0.15,
betas=(0.9, 0.95)
)
tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/cosmo2-tokenizer")
tokenizer.pad_token = tokenizer.eos_token
vocab_size = tokenizer.vocab_size
# Adjusted batch size and sequence length
train_loader = load_cosmopedia_dataset(
batch_size=16, # Set to 16
seq_length=1024, # Kept at 1024
tokenizer=tokenizer
)
import time
t1 = time.time()
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# Set environment variable for memory allocation
import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
train_model(
config,
model,
train_loader,
train_loader,
optimizer=optimizer,
device=device,
num_epochs=1,
eval_freq=1000, # Increase eval frequency to every 500 steps
eval_iter=1000,
start_context="Once Upon a Time far far away in a galaxy",
tokenizer=tokenizer
)
t2 = time.time()
logger.info(f"Time taken for training: {t2 - t1:.2f} seconds")