import argparse import contextlib import copy import json import os import math import time # Persist torch.compile()'d kernels across runs so startup is fast after the first run. os.environ.setdefault( "TORCHINDUCTOR_CACHE_DIR", os.path.expanduser("~/.cache/torch_inductor"), ) import wandb import numpy as np from pathlib import Path from datetime import datetime from typing import List, Optional, Sequence, Tuple, Union from dataclasses import dataclass, asdict from datasets import load_from_disk import torch import torch.distributed as dist from torch.utils.data import Dataset, DataLoader from torch.utils.data.distributed import DistributedSampler from torch.nn.parallel import DistributedDataParallel as DDP from torch.optim.lr_scheduler import LambdaLR from sparseatt.config import Qwen3Config from sparseatt.model import Qwen3ForCausalLM, load_tokenizer import torch._inductor.config as inductor_config try: from torchinfo import summary except ImportError: summary = None torch.set_float32_matmul_precision('high') @dataclass class TrainConfig: """Training configuration loaded from JSON.""" # Model model_path: str = "../qwen3_06b_base" model_dtype: str = "bfloat16" init_from_scratch: bool = False # Data train_path: str = "./data/sample/sample.jsonl" val_path: Optional[str] = None max_seq_length: int = 512 num_workers: int = 8 # Optional data/sampler seed. When omitted, the sampler follows the # global training seed; legacy-compatible configs may set this explicitly. data_seed: Optional[int] = None # Training output_dir: str = "./checkpoints" resume_from_checkpoint: Optional[str] = None # Path to checkpoint dir to resume from force_new_run: bool = False # Set True to start fresh even if a previous checkpoint exists num_epochs: int = 1 max_steps: Optional[int] = None per_device_batch_size: int = 2 gradient_accumulation_steps: int = 8 learning_rate: float = 1e-5 weight_decay: float = 0.01 warmup_steps: int = 100 lr_scheduler: str = "cosine" max_grad_norm: float = 1.0 save_steps: int = 500 eval_steps: int = 500 eval_batches: int = 100 logging_steps: int = 50 seed: int = 42 verbose: bool = False checkpoint_activations: bool = False # Explicit per-rank DataLoader cursor for a deliberate DDP world-size migration. resume_micro_step_override: Optional[int] = None # Distributed strategy: str = "ddp" mixed_precision: str = "bf16" # Wandb wandb_enabled: bool = True wandb_project: str = "qwen3-training" wandb_name: Optional[str] = None wandb_tags: Union[list, None] = None # Architecture overrides (applied on top of the base model config.json) hidden_size: Optional[int] = None intermediate_size: Optional[int] = None num_attention_heads: Optional[int] = None num_key_value_heads: Optional[int] = None num_hidden_layers: Optional[int] = None head_dim: Optional[int] = None attention_dropout: Optional[float] = None rope_theta: Optional[float] = None rope_train_position_scale_set: Optional[List[float]] = None positional: str = "rope" qk_norm: str = "on" qk_gain_fp32: bool = False # Attention mode attention_type: str = "dense" sparse_block_size: int = 128 sparse_num_random_blocks: int = 1 sparse_mask_pool_size: int = 1000 sparse_val_dense_compare: bool = False sparse_layer_indices: Optional[List[int]] = None sparse_mask_cache_dir: Optional[str] = None sparse_mask_seed: int = 0 sparse_mixed_nrb_set: Optional[List[int]] = None dense_mode: bool = False dense_step_intervals: Optional[List[List[int]]] = None @classmethod def from_json(cls, path: str) -> "TrainConfig": with open(path, "r") as f: data = json.load(f) config = cls() # Flatten nested config if "model" in data: config.model_path = data["model"].get("path", config.model_path) config.model_dtype = data["model"].get("dtype", config.model_dtype) config.init_from_scratch = data["model"].get("init_from_scratch", config.init_from_scratch) config.hidden_size = data["model"].get("hidden_size", config.hidden_size) config.intermediate_size = data["model"].get("intermediate_size", config.intermediate_size) config.num_attention_heads = data["model"].get("num_attention_heads", config.num_attention_heads) config.num_key_value_heads = data["model"].get("num_key_value_heads", config.num_key_value_heads) config.num_hidden_layers = data["model"].get("num_hidden_layers", config.num_hidden_layers) config.head_dim = data["model"].get("head_dim", config.head_dim) config.attention_dropout = data["model"].get("attention_dropout", config.attention_dropout) config.rope_theta = data["model"].get("rope_theta", config.rope_theta) config.rope_train_position_scale_set = data["model"].get( "rope_train_position_scale_set", config.rope_train_position_scale_set ) config.positional = data["model"].get("positional", config.positional) config.qk_norm = data["model"].get("qk_norm", config.qk_norm) config.qk_gain_fp32 = data["model"].get("qk_gain_fp32", config.qk_gain_fp32) if "data" in data: config.train_path = data["data"].get("train_path", config.train_path) config.val_path = data["data"].get("val_path", config.val_path) config.max_seq_length = data["data"].get("max_seq_length", config.max_seq_length) config.num_workers = data["data"].get("num_workers", config.num_workers) config.data_seed = data["data"].get("seed", config.data_seed) if "training" in data: t = data["training"] config.output_dir = t.get("output_dir", config.output_dir) config.resume_from_checkpoint = t.get("resume_from_checkpoint", config.resume_from_checkpoint) config.force_new_run = t.get("force_new_run", config.force_new_run) config.num_epochs = t.get("num_epochs", config.num_epochs) config.max_steps = t.get("max_steps", config.max_steps) config.per_device_batch_size = t.get("per_device_batch_size", config.per_device_batch_size) config.gradient_accumulation_steps = t.get("gradient_accumulation_steps", config.gradient_accumulation_steps) config.learning_rate = t.get("learning_rate", config.learning_rate) config.weight_decay = t.get("weight_decay", config.weight_decay) config.warmup_steps = t.get("warmup_steps", config.warmup_steps) config.lr_scheduler = t.get("lr_scheduler", config.lr_scheduler) config.max_grad_norm = t.get("max_grad_norm", config.max_grad_norm) config.save_steps = t.get("save_steps", config.save_steps) config.eval_steps = t.get("eval_steps", config.eval_steps) config.eval_batches = t.get("eval_batches", config.eval_batches) config.logging_steps = t.get("logging_steps", config.logging_steps) config.seed = t.get("seed", config.seed) config.verbose = t.get("verbose", config.verbose) config.checkpoint_activations = t.get("checkpoint_activations", config.checkpoint_activations) config.resume_micro_step_override = t.get("resume_micro_step_override", config.resume_micro_step_override) if "distributed" in data: config.strategy = data["distributed"].get("strategy", config.strategy) config.mixed_precision = data["distributed"].get("mixed_precision", config.mixed_precision) if "wandb" in data: config.wandb_enabled = data["wandb"].get("enabled", config.wandb_enabled) config.wandb_project = data["wandb"].get("project", config.wandb_project) config.wandb_name = data["wandb"].get("name", config.wandb_name) config.wandb_tags = data["wandb"].get("tags", config.wandb_tags) attention_cfg = data.get("attention", {}) config.attention_type = attention_cfg.get("type", data.get("attention_type", config.attention_type)) config.sparse_block_size = attention_cfg.get("sparse_block_size", data.get("sparse_block_size", config.sparse_block_size)) config.sparse_num_random_blocks = attention_cfg.get( "sparse_num_random_blocks", data.get("sparse_num_random_blocks", config.sparse_num_random_blocks), ) config.sparse_mask_pool_size = attention_cfg.get( "sparse_mask_pool_size", data.get("sparse_mask_pool_size", config.sparse_mask_pool_size), ) config.sparse_mask_seed = attention_cfg.get( "sparse_mask_seed", data.get("sparse_mask_seed", config.sparse_mask_seed) ) config.sparse_mixed_nrb_set = attention_cfg.get( "sparse_mixed_nrb_set", data.get("sparse_mixed_nrb_set", config.sparse_mixed_nrb_set), ) config.sparse_val_dense_compare = attention_cfg.get( "sparse_val_dense_compare", attention_cfg.get( "sparse_validate_dense_compare", data.get("sparse_val_dense_compare", config.sparse_val_dense_compare), ), ) config.sparse_layer_indices = attention_cfg.get( "sparse_layer_indices", data.get("sparse_layer_indices", config.sparse_layer_indices), ) config.sparse_mask_cache_dir = attention_cfg.get( "sparse_mask_cache_dir", data.get("sparse_mask_cache_dir", config.sparse_mask_cache_dir), ) config.dense_mode = attention_cfg.get( "dense_mode", data.get("dense_mode", config.dense_mode), ) config.dense_step_intervals = attention_cfg.get( "dense_step_intervals", attention_cfg.get( "dense_mode_step_intervals", data.get("dense_step_intervals", data.get("dense_mode_step_intervals", config.dense_step_intervals)), ), ) return config def _normalize_dense_step_intervals( raw_intervals: Optional[Sequence[Sequence[int]]], ) -> List[Tuple[int, int]]: if raw_intervals in (None, []): return [] if not isinstance(raw_intervals, Sequence) or isinstance(raw_intervals, (str, bytes)): raise ValueError("dense_step_intervals must be a list of [start, end] pairs") normalized: List[Tuple[int, int]] = [] for idx, interval in enumerate(raw_intervals): if not isinstance(interval, Sequence) or isinstance(interval, (str, bytes)) or len(interval) != 2: raise ValueError( f"dense_step_intervals[{idx}] must be a [start, end] pair" ) start, end = interval if not isinstance(start, int) or not isinstance(end, int): raise ValueError( f"dense_step_intervals[{idx}] must contain integer step indices" ) if start < 0 or end < 0: raise ValueError( f"dense_step_intervals[{idx}] must contain non-negative step indices" ) if start > end: raise ValueError( f"dense_step_intervals[{idx}] start ({start}) must be <= end ({end})" ) normalized.append((start, end)) normalized.sort(key=lambda pair: pair[0]) for idx in range(1, len(normalized)): prev_start, prev_end = normalized[idx - 1] start, end = normalized[idx] if start <= prev_end: raise ValueError( "dense_step_intervals must be non-overlapping when endpoints are inclusive; " f"got [{prev_start}, {prev_end}] and [{start}, {end}]" ) return normalized def _step_uses_dense_attention(step: int, dense_step_intervals: Sequence[Tuple[int, int]]) -> bool: if step <= 0: return False return any(start <= step <= end for start, end in dense_step_intervals) def _active_dense_step_interval( step: int, dense_step_intervals: Sequence[Tuple[int, int]], ) -> Optional[Tuple[int, int]]: if step <= 0: return None for start, end in dense_step_intervals: if start <= step <= end: return start, end return None def _set_model_dense_mode(model, enabled: bool) -> None: base_model = _unwrap_base_model(model) base_model.enable_dense_mode(enabled) def _apply_dense_mode_for_step( model, config: TrainConfig, step: int, dense_step_intervals: Sequence[Tuple[int, int]], ) -> bool: if config.attention_type != "sparse_stochastic": return False enabled = config.dense_mode or _step_uses_dense_attention(step, dense_step_intervals) _set_model_dense_mode(model, enabled) return enabled class TextDataset(Dataset): """Simple text dataset for language model training.""" def __init__( self, data_path: str, tokenizer, max_seq_length: int = 512, ): self.tokenizer = tokenizer self.max_seq_length = max_seq_length self.examples = [] # Load data if data_path.endswith(".jsonl"): self._load_jsonl(data_path) else: # Assume it's a directory with HF dataset self._load_hf_dataset(data_path) def _load_jsonl(self, path: str): with open(path, "r") as f: for line in f: data = json.loads(line) text = data.get("text", data.get("content", "")) if text.strip(): self.examples.append(text) def _load_hf_dataset(self, path: str): dataset = load_from_disk(path) for example in dataset: text = example.get("text", example.get("content", "")) if text.strip(): self.examples.append(text) def __len__(self): return len(self.examples) def __getitem__(self, idx): text = self.examples[idx] # Tokenize encodings = self.tokenizer( text, truncation=True, max_length=self.max_seq_length, padding="max_length", return_tensors="pt", ) input_ids = encodings["input_ids"].squeeze(0) attention_mask = encodings["attention_mask"].squeeze(0) # Labels are same as input_ids for causal LM labels = input_ids.clone() # Mask padding tokens in labels labels[attention_mask == 0] = -100 return { "input_ids": input_ids, "attention_mask": attention_mask, "labels": labels, } class PreTokenizedDataset(Dataset): """Dataset for pre-tokenized data (no tokenization at runtime).""" def __init__(self, data_path: str, max_seq_length: int = 2048): """ Load pre-tokenized data. Args: data_path: Path to directory containing tokenized_data.pt or path to HuggingFace dataset directory with tokenized data """ self.data_path = data_path self.max_seq_length = max_seq_length # Check if it's a .pt file or directory if os.path.isfile(data_path) and data_path.endswith(".pt"): # Load from .pt file print(f"Loading pre-tokenized data from {data_path}...") self.data = torch.load(data_path) self.is_hf_dataset = False elif os.path.isdir(data_path): # Check for tokenized_data.pt in directory pt_path = os.path.join(data_path, "tokenized_data.pt") if os.path.exists(pt_path): print(f"Loading pre-tokenized data from {pt_path}...") self.data = torch.load(pt_path) self.is_hf_dataset = False else: # Try loading as HuggingFace dataset print(f"Loading pre-tokenized HF dataset from {data_path}...") self.data = load_from_disk(data_path) # Set format to torch to avoid repeated tensor conversions self.data.set_format(type="torch", columns=["input_ids", "attention_mask"]) self.is_hf_dataset = True else: raise ValueError(f"Invalid data path: {data_path}") # Compute pack factor: how many stored sequences to concatenate per example. # Uses ceil so any max_seq_length >= stored_len works, not just exact multiples. self.pack_factor = 1 self.stored_seq_length = max_seq_length if self.is_hf_dataset and len(self.data) > 0: stored_len = len(self.data[0]["input_ids"]) if stored_len > 0 and stored_len < max_seq_length: self.pack_factor = math.ceil(max_seq_length / stored_len) self.stored_seq_length = stored_len n_packed = len(self.data) // self.pack_factor print(f"Loaded {len(self.data)} pre-tokenized examples " f"(stored_len={self.stored_seq_length}, pack_factor={self.pack_factor}, " f"packed_examples={n_packed})") def __len__(self): return len(self.data) // self.pack_factor def __getitem__(self, idx): if self.is_hf_dataset: if self.pack_factor > 1: # Concatenate pack_factor consecutive stored sequences into one start = idx * self.pack_factor ids_list, mask_list = [], [] for i in range(self.pack_factor): ex = self.data[start + i] ids_list.append(ex["input_ids"]) mask_list.append(ex["attention_mask"]) input_ids = torch.cat(ids_list, dim=0)[:self.max_seq_length] attention_mask = torch.cat(mask_list, dim=0)[:self.max_seq_length] else: example = self.data[idx] # Data is already in tensor format due to set_format() input_ids = example["input_ids"][:self.max_seq_length] attention_mask = example["attention_mask"][:self.max_seq_length] labels = input_ids.clone() # Mask padding tokens in labels labels[attention_mask == 0] = -100 return { "input_ids": input_ids, "attention_mask": attention_mask, "labels": labels, } else: # Data is already in tensor format return self.data[idx][:self.max_seq_length] def get_lr_scheduler(optimizer, config: TrainConfig, num_training_steps: int): """Create learning rate scheduler.""" num_warmup_steps = config.warmup_steps if config.lr_scheduler == "cosine": def lr_lambda(current_step): if current_step < num_warmup_steps: return float(current_step) / float(max(1, num_warmup_steps)) progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress))) elif config.lr_scheduler == "linear": def lr_lambda(current_step): if current_step < num_warmup_steps: return float(current_step) / float(max(1, num_warmup_steps)) return max(0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))) else: # constant def lr_lambda(current_step): if current_step < num_warmup_steps: return float(current_step) / float(max(1, num_warmup_steps)) return 1.0 return LambdaLR(optimizer, lr_lambda) def setup_distributed(): """Setup distributed training.""" if "RANK" in os.environ: rank = int(os.environ["RANK"]) local_rank = int(os.environ["LOCAL_RANK"]) world_size = int(os.environ["WORLD_SIZE"]) dist.init_process_group(backend="nccl") torch.cuda.set_device(local_rank) return rank, local_rank, world_size else: return 0, 0, 1 def is_main_process(rank: int) -> bool: return rank == 0 def _unwrap_model(model): return model.module if hasattr(model, "module") else model def _unwrap_base_model(model): unwrapped = _unwrap_model(model) return unwrapped._orig_mod if hasattr(unwrapped, "_orig_mod") else unwrapped def _model_attention_mask(attention_mask): """Map a batch attention mask to what the model should receive. Returns ``None`` for fully-unpadded (packed) batches so dense / dense_mode attention layers take SDPA's ``is_causal`` fast path (keeps the Flash backend eligible and applies the causal upper-triangle skip). Returns the mask unchanged when padding is present so the explicit additive mask is built. The padding check runs here in eager mode — deliberately outside the compiled model — so it never introduces a graph break / recompile in the hot forward. Loss masking is unaffected: labels already carry ``-100`` for padded positions, so dropping an all-ones attention mask does not change the loss. """ if attention_mask is None: return None return attention_mask if bool((attention_mask == 0).any()) else None def _sample_sparse_indices(model, config: TrainConfig, mask_rng: Optional[np.random.RandomState]): if config.attention_type != "sparse_stochastic": return None if mask_rng is None: raise ValueError("mask_rng must be provided for sparse_stochastic attention") base_model = _unwrap_base_model(model) layers = base_model.model.layers pool_size = int(base_model.config.sparse_mask_pool_size) indices = mask_rng.randint(0, pool_size, size=len(layers)) # Pre-select BlockMask objects outside the compiled graph so Dynamo never # sees list-indexing with a dynamic integer. flex_attention's HOP # registration guards only on BlockMask structural metadata (Q_LEN, # KV_LEN, BLOCK_SIZE) — not on which random blocks are chosen — so # different masks across steps cause zero recompilations. return [ layer.self_attn.pick_mask(int(idx)) if layer.is_sparse else None for layer, idx in zip(layers, indices) ] def _build_dense_eval_model(train_model, device: torch.device, dtype: torch.dtype): train_base = _unwrap_base_model(train_model) dense_cfg = copy.deepcopy(train_base.config) dense_cfg.attention_type = "dense" return Qwen3ForCausalLM(dense_cfg).to(device=device, dtype=dtype) def _sync_dense_eval_model_weights(train_model, dense_eval_model): train_base = _unwrap_base_model(train_model) missing, unexpected = dense_eval_model.load_state_dict(train_base.state_dict(), strict=False) if missing or unexpected: print(f"[sync_dense_eval] missing={len(missing)} unexpected={len(unexpected)}") if missing: print(" missing[:10]:", missing[:10]) if unexpected: print(" unexpected[:10]:", unexpected[:10]) def evaluate( model, eval_loader, device, config: TrainConfig, rank: int, mask_rng: Optional[np.random.RandomState] = None, use_sparse_indices: bool = True, dense_step_intervals: Sequence[Tuple[int, int]] = (), step_for_mode: Optional[int] = None, ): """Evaluate model on validation set.""" model.eval() total_loss = 0.0 total_tokens = 0 num_batches = 0 # Divide the target by world_size so the work is distributed evenly world_size = dist.get_world_size() if dist.is_initialized() else 1 # Use math.ceil to ensure you don't round down to 0 on small batch settings max_eval_batches = math.ceil( (config.eval_batches * config.gradient_accumulation_steps) / world_size ) autocast_dtype = torch.bfloat16 if config.mixed_precision == "bf16" else torch.float16 dense_active = False if step_for_mode is not None: dense_active = _apply_dense_mode_for_step(model, config, step_for_mode, dense_step_intervals) with torch.no_grad(): for batch in eval_loader: if num_batches >= max_eval_batches: break input_ids = batch["input_ids"].to(device) attention_mask = batch["attention_mask"].to(device) labels = batch["labels"].to(device) sparse_indices = ( _sample_sparse_indices(model, config, mask_rng) if use_sparse_indices and not dense_active else None ) with torch.amp.autocast("cuda", dtype=autocast_dtype): outputs = model( input_ids=input_ids, attention_mask=_model_attention_mask(attention_mask), labels=labels, sparse_indices=sparse_indices, ) loss = outputs["loss"] total_loss += loss.item() total_tokens += attention_mask.sum().item() num_batches += 1 avg_loss = total_loss / num_batches if num_batches > 0 else 0.0 # evaluate() is called only on rank 0 (gated by is_main_process in the train loop). # All-reduce across ranks is therefore unreachable and has been removed to avoid # potential deadlocks if the gating is ever adjusted. # Calculate perplexity perplexity = math.exp(avg_loss) if avg_loss < 100 else float('inf') model.train() return avg_loss, perplexity def run_validation_step( model, val_loader, device, config: TrainConfig, rank: int, mask_rng: Optional[np.random.RandomState], global_step: int, dense_eval_model=None, dense_step_intervals: Sequence[Tuple[int, int]] = (), ): """Run validation across all ranks to prevent DDP deadlocks, logging only on rank 0.""" active_dense_interval = _active_dense_step_interval(global_step, dense_step_intervals) validation_uses_dense_attention = config.dense_mode or active_dense_interval is not None # CRITICAL: Every GPU must run this so they stay synchronized in the timeline eval_loss, eval_perplexity = evaluate( model, val_loader, device, config, rank, mask_rng, dense_step_intervals=dense_step_intervals, step_for_mode=global_step, ) dense_eval_loss = None dense_eval_perplexity = None if dense_eval_model is not None: _sync_dense_eval_model_weights(model, dense_eval_model) dense_eval_loss, dense_eval_perplexity = evaluate( dense_eval_model, val_loader, device, config, rank, mask_rng=None, use_sparse_indices=False, ) # GATING ONLY THE OUTPUTS: Only rank 0 prints and logs to WandB if rank == 0: print(f"Validation at step {global_step}: Loss = {eval_loss:.4f}, Perplexity = {eval_perplexity:.2f}") log_dict = { "eval/loss": eval_loss, "eval/perplexity": eval_perplexity, "eval/uses_dense_attention": int(validation_uses_dense_attention), "eval/uses_sparse_attention": int(not validation_uses_dense_attention), } if active_dense_interval is not None: interval_start, interval_end = active_dense_interval log_dict.update({ "eval/dense_interval_active": 1, "eval/dense_interval_start": interval_start, "eval/dense_interval_end": interval_end, }) else: log_dict.update({ "eval/dense_interval_active": 0, "eval/dense_interval_start": -1, "eval/dense_interval_end": -1, }) if dense_eval_model is not None: print( f"Dense validation compare at step {global_step}: " f"Loss = {dense_eval_loss:.4f}, Perplexity = {dense_eval_perplexity:.2f}" ) log_dict.update({ "eval_dense_compare/loss": dense_eval_loss, "eval_dense_compare/perplexity": dense_eval_perplexity, }) if config.wandb_enabled: wandb.log(log_dict, step=global_step) return log_dict return None def load_checkpoint(checkpoint_path: str, model, optimizer, scheduler, device, rank: int, config: TrainConfig | None = None): """Load model checkpoint and training state. Returns ------- (global_step, resume_epoch, resume_micro_step) resume_epoch – epoch index to start the training loop from resume_micro_step – number of micro-batches already consumed in that epoch; the loop should skip this many batches before training. """ if is_main_process(rank): print(f"Loading checkpoint from {checkpoint_path}...") checkpoint_dir = Path(checkpoint_path) # Load model weights model_path = checkpoint_dir / "model.pt" if not model_path.exists(): raise FileNotFoundError(f"Model checkpoint not found: {model_path}") model_state = torch.load(model_path, map_location=device, weights_only=True) # Normalise checkpoint keys: strip "_orig_mod." prefix if present (checkpoint # was saved from a torch.compile()-wrapped model). if any(k.startswith("_orig_mod.") for k in model_state): model_state = {k.removeprefix("_orig_mod."): v for k, v in model_state.items()} # Drop _rt_pool entries — they are non-learned buffers rebuilt each run and # their shape depends on sparse_mask_pool_size which can differ across configs. model_state = {k: v for k, v in model_state.items() if not k.endswith("._rt_pool")} # Unwrap DDP, then torch.compile's OptimizedModule, so plain keys always match. target = model.module if hasattr(model, "module") else model if hasattr(target, "_orig_mod"): target = target._orig_mod target.load_state_dict(model_state) # Load training state (optimizer, scheduler, step, epoch, micro_step) training_state_path = checkpoint_dir / "training_state.pt" if not training_state_path.exists(): raise FileNotFoundError(f"Training state not found: {training_state_path}") training_state = torch.load(training_state_path, map_location=device) optimizer.load_state_dict(training_state["optimizer"]) scheduler.load_state_dict(training_state["scheduler"]) # Override both optimizer and scheduler LR state only when the configured # schedule base really changed. For an exact resume, the config contains # the original peak LR while the checkpoint optimizer contains the current # (scheduled) LR. Unconditionally assigning the peak LR here would create # a one-step spike before scheduler.step() restores the schedule. if config is not None: new_lr = config.learning_rate checkpoint_base_lrs = list(getattr(scheduler, "base_lrs", [])) same_schedule_base = bool(checkpoint_base_lrs) and all( math.isclose(base_lr, new_lr, rel_tol=1e-12, abs_tol=0.0) for base_lr in checkpoint_base_lrs ) if same_schedule_base: if is_main_process(rank): current_lrs = [pg["lr"] for pg in optimizer.param_groups] print( "Preserving checkpoint LR state for exact resume: " f"current={current_lrs}, schedule_base={checkpoint_base_lrs}" ) else: optimizer.defaults["lr"] = new_lr for pg in optimizer.param_groups: pg["lr"] = new_lr pg["initial_lr"] = new_lr if hasattr(scheduler, "base_lrs"): scheduler.base_lrs = [new_lr for _ in scheduler.base_lrs] if hasattr(scheduler, "_last_lr"): scheduler._last_lr = [new_lr for _ in scheduler._last_lr] if is_main_process(rank): print(f"Overriding checkpoint LR with config LR: {new_lr}") step = training_state["step"] resume_epoch = training_state.get("epoch", 0) resume_micro_step = training_state.get("micro_step", 0) # Restore per-rank RNG states so stochastic ops are exactly reproducible. rng_path = checkpoint_dir / f"rng_state_rank{rank}.pt" if rng_path.exists(): # weights_only=False is required because the file contains numpy arrays # (torch CPU RNG state and numpy RandomState). The file is always written # by this codebase so the source is trusted. rng_state = torch.load(rng_path, map_location="cpu", weights_only=False) torch.set_rng_state(rng_state["torch_cpu"]) if torch.cuda.is_available() and rng_state.get("torch_cuda") is not None: torch.cuda.set_rng_state(rng_state["torch_cuda"]) np.random.set_state(rng_state["numpy"]) # mask_rng is a local RandomState; its state is returned so the caller # can apply it after construction. else: if is_main_process(rank): print(f"Warning: no RNG state file found for rank {rank}; RNG will not be restored.") rng_state = None if is_main_process(rank): print(f"Resumed from checkpoint at step {step} (epoch {resume_epoch}, micro_step {resume_micro_step})") return step, resume_epoch, resume_micro_step, rng_state def save_checkpoint( model, optimizer, scheduler, step: int, config: TrainConfig, rank: int, epoch: int = 0, micro_step: int = 0, mask_rng: Optional[np.random.RandomState] = None, loader_generator: Optional[torch.Generator] = None, sampler_seed: Optional[int] = None, ): """Save model checkpoint, dataloader position, and per-rank RNG states.""" checkpoint_dir = Path(config.output_dir) / f"checkpoint-{step}" checkpoint_dir.mkdir(parents=True, exist_ok=True) # Every rank saves its own RNG state and DataLoader generator state for exact resume. rng_state = { "torch_cpu": torch.get_rng_state(), "torch_cuda": torch.cuda.get_rng_state() if torch.cuda.is_available() else None, "numpy": np.random.get_state(), "mask_rng": mask_rng.get_state() if mask_rng is not None else None, } torch.save(rng_state, checkpoint_dir / f"rng_state_rank{rank}.pt") if loader_generator is not None: torch.save(loader_generator.get_state(), checkpoint_dir / f"loader_gen_state_rank{rank}.pt") if not is_main_process(rank): return # Get model state dict (handle DDP) model_to_save = model.module if hasattr(model, "module") else model # Save model torch.save(model_to_save.state_dict(), checkpoint_dir / "model.pt") # Save optimizer, scheduler, and dataloader position torch.save({ "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict(), "step": step, "epoch": epoch, "micro_step": micro_step, "sampler_seed": sampler_seed, }, checkpoint_dir / "training_state.pt") # Save config with open(checkpoint_dir / "config.json", "w") as f: json.dump(vars(model_to_save.config), f, indent=2, default=str) print(f"Saved checkpoint to {checkpoint_dir}") def find_latest_checkpoint(output_dir: str) -> Optional[str]: """Return the path of the most recent checkpoint inside output_dir/run_*/, or None.""" output_path = Path(output_dir) if not output_path.exists(): return None required_files = ("model.pt", "training_state.pt", "config.json") run_dirs = sorted(output_path.glob("run_*"), reverse=True) # newest name first for run_dir in run_dirs: checkpoints = sorted( ( checkpoint for checkpoint in run_dir.glob("checkpoint-*") if all( (checkpoint / filename).is_file() and (checkpoint / filename).stat().st_size > 0 for filename in required_files ) ), key=lambda p: int(p.name.split("-")[-1]), reverse=True, ) if checkpoints: return str(checkpoints[0]) return None def resolve_sampler_seed(config: TrainConfig) -> int: """Choose the sampler seed without changing legacy resume order. Checkpoints written before the sampler-seed refactor have no marker and used DistributedSampler's default seed (0). New runs use config.seed and persist that choice in training_state.pt for subsequent exact resumes. """ if not config.resume_from_checkpoint: return int(config.data_seed if config.data_seed is not None else config.seed) state_path = Path(config.resume_from_checkpoint) / "training_state.pt" try: state = torch.load(state_path, map_location="cpu", weights_only=True) value = state.get("sampler_seed") if value is not None: return int(value) except Exception: pass return int(config.data_seed if config.data_seed is not None else 0) def train(config: TrainConfig): """Main training function.""" # Setup distributed rank, local_rank, world_size = setup_distributed() device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu") # Set seed torch.manual_seed(config.seed + rank) mask_rng = np.random.RandomState(config.seed + 1 + rank) # mask_rng state is restored after load_checkpoint (below) once resume_mask_rng_state is known. if config.attention_type not in {"dense", "sparse_stochastic"}: raise ValueError( f"Unsupported attention_type={config.attention_type}. Expected 'dense' or 'sparse_stochastic'." ) if config.dense_mode and config.attention_type != "sparse_stochastic": raise ValueError( "dense_mode=True requires attention_type='sparse_stochastic' so the sparse layers exist to switch over." ) dense_step_intervals = _normalize_dense_step_intervals(config.dense_step_intervals) config.dense_step_intervals = [list(interval) for interval in dense_step_intervals] or None if dense_step_intervals and config.attention_type != "sparse_stochastic": raise ValueError( "dense_step_intervals requires attention_type='sparse_stochastic'." ) if dense_step_intervals and config.dense_mode: raise ValueError( "dense_step_intervals cannot be combined with dense_mode=True; use one dense fallback mechanism." ) # Auto-detect latest checkpoint to resume from if not explicitly specified if not config.resume_from_checkpoint and not config.force_new_run: latest_ckpt = find_latest_checkpoint(config.output_dir) if latest_ckpt is not None: config.resume_from_checkpoint = latest_ckpt if is_main_process(rank): print(f"Auto-resuming from latest checkpoint: {latest_ckpt}") # Generate unique run ID and update output directory if config.resume_from_checkpoint and not config.force_new_run: # Truly resuming the same run — reuse its output directory so wandb # and optimizer state continue from where they left off. config.output_dir = str(Path(config.resume_from_checkpoint).parent) if is_main_process(rank): print(f"Resuming: Using existing output directory {config.output_dir}") else: # Either a completely fresh run or a fork (force_new_run=True with a # resume_from_checkpoint for weights only) — always create a new dir # so a new wandb run is created instead of hijacking the old one. run_id = datetime.now().strftime("%Y%m%d_%H%M%S") config.output_dir = os.path.join(config.output_dir, f"run_{run_id}") if is_main_process(rank): print(f"New run: Checkpoints will be saved to {config.output_dir}") Path(config.output_dir).mkdir(parents=True, exist_ok=True) # Persist wandb run ID early (before heavy setup) so if the process crashes # during model/dataset loading, recovery will continue in the same wandb run. # Actual wandb.init() is deferred to just before training loop to ensure # wall-clock timing only measures training time, not setup time. wandb_run_id_path = None prior_wandb_id = None if config.wandb_enabled and is_main_process(rank): wandb_run_id_path = Path(config.output_dir) / "wandb_run_id.txt" if wandb_run_id_path.exists(): prior_wandb_id = wandb_run_id_path.read_text().strip() if is_main_process(rank): print("Running config:") print(json.dumps(asdict(config), indent=2, default=str)) # Load tokenizer if is_main_process(rank): print(f"Loading tokenizer from {config.model_path}...") tokenizer = load_tokenizer(config.model_path) # Set padding token if not set if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token # Load model dtype = torch.bfloat16 if config.model_dtype == "bfloat16" else torch.float16 attention_overrides = { "attention_type": config.attention_type, "positional": config.positional, "qk_norm": config.qk_norm, "qk_gain_fp32": config.qk_gain_fp32, "sparse_block_size": int(config.sparse_block_size), "sparse_num_random_blocks": int(config.sparse_num_random_blocks), "sparse_mask_pool_size": int(config.sparse_mask_pool_size), "sparse_mask_seed": int(config.sparse_mask_seed), "sparse_mixed_nrb_set": config.sparse_mixed_nrb_set, "sparse_max_seq_length": int(config.max_seq_length), "sparse_layer_indices": config.sparse_layer_indices, "sparse_mask_cache_dir": config.sparse_mask_cache_dir, } # Add any architecture overrides specified in the train config arch_override_keys = [ "hidden_size", "intermediate_size", "num_attention_heads", "num_key_value_heads", "num_hidden_layers", "head_dim", "attention_dropout", "rope_theta", "qk_norm", "qk_gain_fp32", "rope_train_position_scale_set", ] for key in arch_override_keys: val = getattr(config, key, None) if val is not None: attention_overrides[key] = val # Determine model loading strategy if config.resume_from_checkpoint: # When resuming, load config from checkpoint if is_main_process(rank): print(f"Resuming from checkpoint: {config.resume_from_checkpoint}") checkpoint_config_path = os.path.join(config.resume_from_checkpoint, "config.json") model_config = Qwen3Config.from_json(checkpoint_config_path) model_config.checkpoint_activations = config.checkpoint_activations for key, value in attention_overrides.items(): if hasattr(model_config, key): setattr(model_config, key, value) model = Qwen3ForCausalLM(model_config) model = model.to(device=device, dtype=dtype) # Model weights will be loaded later along with optimizer/scheduler elif config.init_from_scratch: if is_main_process(rank): print(f"Initializing model from scratch (random weights)...") # Load config from pretrained path but initialize with random weights config_path = os.path.join(config.model_path, "config.json") model_config = Qwen3Config.from_json(config_path) model_config.checkpoint_activations = config.checkpoint_activations for key, value in attention_overrides.items(): if hasattr(model_config, key): setattr(model_config, key, value) model = Qwen3ForCausalLM(model_config) model = model.to(device=device, dtype=dtype) else: if is_main_process(rank): print(f"Loading pretrained model from {config.model_path}...") model = Qwen3ForCausalLM.from_pretrained( config.model_path, device=device, dtype=dtype, config_overrides=attention_overrides, ) model.config.checkpoint_activations = config.checkpoint_activations model.model.checkpoint_activations = config.checkpoint_activations if config.dense_mode: model.enable_dense_mode() if is_main_process(rank): print("dense_mode=True: sparse attention layers will use standard causal SDPA.") # ``model.to(..., dtype=bf16)`` above intentionally converts the whole # model. Restore only Q/K RMSNorm gains to FP32 for this experiment; their # forward method casts the value back to the activation dtype. if config.qk_gain_fp32: qk_gain_params = [] for module in model.modules(): if getattr(module, "weight_fp32", False): module.weight.data = module.weight.data.float() qk_gain_params.append(module.weight) if len(qk_gain_params) != 56: raise RuntimeError(f"expected 56 Q/K gain tensors, found {len(qk_gain_params)}") else: qk_gain_params = [] # Print model architecture BEFORE torch.compile so torchinfo hooks don't # survive into the compiled region (which would cause recompile storms). if is_main_process(rank): print() print("\n" + "=" * 80) print("MODEL ARCHITECTURE") print("=" * 80) print(f"Model config:") print(f" vocab_size: {model.config.vocab_size}") print(f" hidden_size: {model.config.hidden_size}") print(f" num_hidden_layers: {model.config.num_hidden_layers}") print(f" num_attention_heads: {model.config.num_attention_heads}") print(f" num_key_value_heads: {model.config.num_key_value_heads}") print(f" head_dim: {model.config.head_dim}") print(f" intermediate_size: {model.config.intermediate_size}") print(f" max_position_embeddings: {model.config.max_position_embeddings}") print(f" rope_theta: {model.config.rope_theta}") print(f" rope_train_position_scale_set: {model.config.rope_train_position_scale_set}") print(f" attention_type: {model.config.attention_type}") if model.config.attention_type == "sparse_stochastic": print(f" sparse_block_size: {model.config.sparse_block_size}") print(f" sparse_num_random_blocks: {model.config.sparse_num_random_blocks}") print(f" sparse_mask_pool_size: {model.config.sparse_mask_pool_size}") print(f" sparse_max_seq_length: {model.config.sparse_max_seq_length}") # Count parameters total_params = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f"\nParameter count:") print(f" Total parameters: {total_params:,}") print(f" Trainable parameters: {trainable_params:,}") print(f" Non-trainable parameters: {total_params - trainable_params:,}") print(f" Model size (BF16): {total_params * 2 / 1e9:.2f} GB") # Print detailed layer summary if torchinfo is available if summary is not None and config.verbose: print("\n" + "-" * 80) print("LAYER-BY-LAYER SUMMARY") print("-" * 80) try: summary( model, input_size=(config.per_device_batch_size, config.max_seq_length), dtypes=[torch.long], col_names=["input_size", "output_size", "num_params", "trainable"], depth=3, verbose=config.verbose, ) except Exception as e: print(f"Could not generate detailed summary: {e}") if getattr(model.config, "tie_word_embeddings", False): print(f" NOTE: torchinfo double-counts tied embed_tokens/lm_head weight " f"({model.config.vocab_size * model.config.hidden_size:,} params). " f"True unique parameter count: {total_params:,} ({total_params/1e6:.1f}M)") print("=" * 80 + "\n") if config.attention_type == "sparse_stochastic": try: backends = getattr(inductor_config, "max_autotune_gemm_backends", "") if isinstance(backends, str): backend_set = {token.strip() for token in backends.split(",") if token.strip()} if "ATEN" not in backend_set: backend_set.add("ATEN") inductor_config.max_autotune_gemm_backends = ",".join(sorted(backend_set)) except Exception as exc: if is_main_process(rank): print(f"Warning: could not set sparse compile backend fallback: {exc}") if os.environ.get("SPARSEATT_DISABLE_COMPILE", "0") != "1": model = torch.compile(model) elif is_main_process(rank): print("torch.compile disabled for this run") dense_eval_model = None if is_main_process(rank) and config.attention_type == "sparse_stochastic" and config.sparse_val_dense_compare: dense_eval_model = _build_dense_eval_model(model, device, dtype) print("Sparse validation compare enabled: running both sparse and dense validation passes") # Wrap with DDP if world_size > 1: # broadcast_buffers=False: _rt_pool is int16 (NCCL unsupported) and is # rebuilt independently per rank each run, so broadcasting it is wrong anyway. model = DDP(model, device_ids=[local_rank], output_device=local_rank, broadcast_buffers=False) # Create dataset and dataloader if is_main_process(rank): print(f"Loading dataset from {config.train_path}...") # Check if data is pre-tokenized is_pretokenized = False if os.path.isdir(config.train_path): # Check for tokenized_data.pt or HF dataset with tokenized columns if os.path.exists(os.path.join(config.train_path, "tokenized_data.pt")): is_pretokenized = True elif os.path.exists(os.path.join(config.train_path, "dataset_info.json")): # HF dataset - check if it has tokenized columns try: dataset_test = load_from_disk(config.train_path) if "input_ids" in dataset_test.column_names: is_pretokenized = True except: pass elif config.train_path.endswith(".pt"): is_pretokenized = True if is_pretokenized: if is_main_process(rank): print("Using pre-tokenized dataset (no runtime tokenization)") full_dataset = PreTokenizedDataset(config.train_path, max_seq_length=config.max_seq_length) else: if is_main_process(rank): print("Using raw text dataset (tokenizing at runtime)") full_dataset = TextDataset( config.train_path, tokenizer, max_seq_length=config.max_seq_length, ) # Do 95/5 train/val split train_size = int(0.95 * len(full_dataset)) val_size = len(full_dataset) - train_size # Set seed for reproducible split generator = torch.Generator().manual_seed(config.seed) train_dataset, val_dataset = torch.utils.data.random_split( full_dataset, [train_size, val_size], generator=generator ) if is_main_process(rank): print(f"Dataset split: {train_size} train, {val_size} validation") # Create train dataloader # Use a dedicated generator (seeded deterministically) for the DataLoader shuffle # so that its RNG is independent from the global training RNG and can be # saved/restored cleanly across checkpoints. loader_generator = torch.Generator() loader_generator.manual_seed(config.seed + rank) sampler_seed = resolve_sampler_seed(config) if is_main_process(rank): source = "checkpoint marker" if config.resume_from_checkpoint and sampler_seed != 0 else ("legacy checkpoint fallback" if config.resume_from_checkpoint else "run seed") print(f"DistributedSampler seed: {sampler_seed} ({source})") if world_size > 1: # Tie the sampler permutation to the same run seed used for model # initialization and the train/validation split. This makes fixed-seed # Dense/Mixed/ALiBi runs consume the same data order when their DDP # world size and loader geometry match. train_sampler = DistributedSampler( train_dataset, num_replicas=world_size, rank=rank, shuffle=True, seed=sampler_seed, ) else: train_sampler = None train_loader = DataLoader( train_dataset, batch_size=config.per_device_batch_size, sampler=train_sampler, shuffle=(train_sampler is None), generator=loader_generator if train_sampler is None else None, num_workers=config.num_workers, pin_memory=True, drop_last=True, persistent_workers=True if config.num_workers > 0 else False, prefetch_factor=2 if config.num_workers > 0 else None, ) # Create validation dataloader if world_size > 1: val_sampler = DistributedSampler( val_dataset, num_replicas=world_size, rank=rank, shuffle=False, seed=sampler_seed, ) else: val_sampler = None val_loader = DataLoader( val_dataset, batch_size=config.per_device_batch_size, sampler=val_sampler, shuffle=False, num_workers=config.num_workers, pin_memory=True, drop_last=False, persistent_workers=True if config.num_workers > 0 else False, prefetch_factor=2 if config.num_workers > 0 else None, ) # Calculate training steps # len(train_loader) is per-rank micro-batches/epoch (DistributedSampler divides data; # drop_last=True in the loader already floors to full batches). micro_batches_per_epoch = len(train_loader) num_update_steps_per_epoch = micro_batches_per_epoch // config.gradient_accumulation_steps num_training_steps = config.max_steps or (num_update_steps_per_epoch * config.num_epochs) effective_batch_size = config.per_device_batch_size * world_size * config.gradient_accumulation_steps tokens_per_step = effective_batch_size * config.max_seq_length if is_main_process(rank): print(f"Training examples (packed, train split): {len(train_dataset)}") print(f"Validation examples (packed, val split): {len(val_dataset)}") print(f"Per-device batch size: {config.per_device_batch_size}") print(f"World size (GPUs): {world_size}") print(f"Gradient accumulation steps: {config.gradient_accumulation_steps}") print(f"Effective batch size (seqs/step): {effective_batch_size}") print(f"Sequence length: {config.max_seq_length}") print(f"Tokens per optimizer step: {tokens_per_step:,}") print(f"Micro-batches per epoch (per rank): {micro_batches_per_epoch}") print(f"Optimizer steps per epoch: {num_update_steps_per_epoch}") print(f"Total optimizer steps: {num_training_steps}" + (f" (capped by max_steps; {num_training_steps/num_update_steps_per_epoch:.2f} epochs)" if config.max_steps else "")) # Create optimizer normal_params = [p for p in model.parameters() if all(p is not q for q in qk_gain_params)] optimizer = torch.optim.AdamW( ([{"params": normal_params}, {"params": qk_gain_params}] if qk_gain_params else normal_params), lr=config.learning_rate, weight_decay=config.weight_decay, betas=(0.9, 0.95), fused=True ) # Create scheduler scheduler = get_lr_scheduler(optimizer, config, num_training_steps) # Load checkpoint if resuming global_step = 0 resume_epoch = 0 resume_micro_step = 0 resume_mask_rng_state = None if config.resume_from_checkpoint: global_step, resume_epoch, resume_micro_step, _rng_state = load_checkpoint( config.resume_from_checkpoint, model, optimizer, scheduler, device, rank, config=config, ) if _rng_state is not None: resume_mask_rng_state = _rng_state.get("mask_rng") if resume_mask_rng_state is not None: mask_rng.set_state(resume_mask_rng_state) if config.resume_micro_step_override is not None: if config.resume_micro_step_override < 0: raise ValueError("resume_micro_step_override must be non-negative") if is_main_process(rank): print(f"Overriding resume micro-step {resume_micro_step} -> {config.resume_micro_step_override} for DDP migration") resume_micro_step = config.resume_micro_step_override if is_main_process(rank): print(f"Resuming training from step {global_step}") # Mixed precision scaler = torch.amp.GradScaler("cuda") if config.mixed_precision == "fp16" else None autocast_dtype = torch.bfloat16 if config.mixed_precision == "bf16" else torch.float16 # Initialize wandb just before training to ensure wall-clock timing # only measures actual training time (not mask creation, model loading, etc.) if config.wandb_enabled and is_main_process(rank): if prior_wandb_id: print(f"Resuming wandb run {prior_wandb_id}") wandb.init( project=config.wandb_project, name=config.wandb_name, tags=config.wandb_tags, config=vars(config), id=prior_wandb_id, resume="must", ) else: wandb.init( project=config.wandb_project, name=config.wandb_name, tags=config.wandb_tags, config=vars(config), ) wandb_run_id_path.write_text(wandb.run.id) # Training loop model.train() total_loss = 0.0 start_time = time.time() last_log_time = start_time last_log_step = global_step _resume_initial_global_step = global_step if is_main_process(rank): print("Starting training...") # Restore loader_generator state from checkpoint so the DataLoader shuffle # on resume is identical to what the original run would have produced. resume_loader_gen_state = None if config.resume_from_checkpoint: _gen_path = Path(config.resume_from_checkpoint) / f"loader_gen_state_rank{rank}.pt" if _gen_path.exists(): resume_loader_gen_state = torch.load(_gen_path, map_location="cpu", weights_only=True) loader_generator.set_state(resume_loader_gen_state) else: if is_main_process(rank): print("Warning: no loader_gen_state file found; DataLoader shuffle order on resume may differ from original run.") # On resume, run one validation pass at the restored global_step before any # new training work, so the baseline metric is logged at the correct step. if config.resume_from_checkpoint and global_step > 0: run_validation_step( model=model, val_loader=val_loader, device=device, config=config, rank=rank, mask_rng=mask_rng, global_step=global_step, dense_eval_model=dense_eval_model, dense_step_intervals=dense_step_intervals, ) last_dense_interval_mode = None for epoch in range(config.num_epochs): # Skip epochs that were fully completed before the checkpoint. if epoch < resume_epoch: continue if train_sampler is not None: train_sampler.set_epoch(epoch) # How many micro-batches to skip at the start of the resume epoch. skip_micro_batches = resume_micro_step if epoch == resume_epoch else 0 for step, batch in enumerate(train_loader): if step < skip_micro_batches: # Fast-forward: advance the DataLoader without training. continue active_step = global_step + 1 dense_interval_active = _apply_dense_mode_for_step( model, config, active_step, dense_step_intervals, ) if dense_interval_active != last_dense_interval_mode: if is_main_process(rank): mode_name = "dense" if dense_interval_active else "sparse_stochastic" source = "dense step interval" if dense_interval_active else "configured sparse attention" print(f"Switching attention mode for step {active_step}: {mode_name} ({source})") last_dense_interval_mode = dense_interval_active # Move batch to device with non-blocking transfer input_ids = batch["input_ids"].to(device, non_blocking=True) attention_mask = batch["attention_mask"].to(device, non_blocking=True) labels = batch["labels"].to(device, non_blocking=True) sparse_indices = None if dense_interval_active else _sample_sparse_indices(model, config, mask_rng) # Forward pass with autocast with torch.amp.autocast("cuda", dtype=autocast_dtype): outputs = model( input_ids=input_ids, attention_mask=_model_attention_mask(attention_mask), labels=labels, sparse_indices=sparse_indices, ) loss = outputs["loss"] / config.gradient_accumulation_steps # Backward pass — suppress DDP gradient sync on non-final accumulation steps # (avoids 7/8 wasted all-reduces when gradient_accumulation_steps=8) is_last_accum_step = (step + 1) % config.gradient_accumulation_steps == 0 sync_ctx = contextlib.nullcontext() if is_last_accum_step else ( model.no_sync() if hasattr(model, "no_sync") else contextlib.nullcontext() ) with sync_ctx: if scaler is not None: scaler.scale(loss).backward() else: loss.backward() # Accumulate UNSCALED loss for logging total_loss += outputs["loss"].item() # Gradient accumulation if (step + 1) % config.gradient_accumulation_steps == 0: # Gradient clipping if scaler is not None: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm) # Optimizer step if scaler is not None: scaler.step(optimizer) scaler.update() else: optimizer.step() scheduler.step() optimizer.zero_grad(set_to_none=True) global_step += 1 # Logging if global_step % config.logging_steps == 0 and is_main_process(rank): avg_loss = total_loss / (config.logging_steps * config.gradient_accumulation_steps) lr = scheduler.get_last_lr()[0] # Calculate speed metrics current_time = time.time() # Instantaneous throughput (since last log) time_since_last_log = current_time - last_log_time steps_since_last_log = global_step - last_log_step steps_per_sec = steps_since_last_log / time_since_last_log if time_since_last_log > 0 else 0 samples_since_last_log = steps_since_last_log * config.gradient_accumulation_steps * config.per_device_batch_size * world_size samples_per_sec = samples_since_last_log / time_since_last_log if time_since_last_log > 0 else 0 ms_per_step = (time_since_last_log / steps_since_last_log * 1000) if steps_since_last_log > 0 else 0 # Update last log tracking last_log_time = current_time last_log_step = global_step print(f"Step {global_step}/{num_training_steps} | Loss: {avg_loss:.4f} | LR: {lr:.2e} | " f"{ms_per_step:.1f}ms/step | {samples_per_sec:.1f} samples/s") if config.wandb_enabled: wandb.log({ "train/loss": avg_loss, "train/learning_rate": lr, "train/epoch": global_step / num_update_steps_per_epoch, "train/global_step": global_step, "speed/steps_per_sec": steps_per_sec, "speed/samples_per_sec": samples_per_sec, "speed/ms_per_step": ms_per_step, }, step=global_step) total_loss = 0.0 # Evaluation if global_step % config.eval_steps == 0: run_validation_step( model=model, val_loader=val_loader, device=device, config=config, rank=rank, mask_rng=mask_rng, global_step=global_step, dense_eval_model=dense_eval_model, dense_step_intervals=dense_step_intervals, ) # Save checkpoint if global_step % config.save_steps == 0: save_checkpoint( model, optimizer, scheduler, global_step, config, rank, epoch=epoch, micro_step=step + 1, mask_rng=mask_rng, loader_generator=loader_generator, sampler_seed=sampler_seed, ) # Check max steps if config.max_steps and global_step >= config.max_steps: break if config.max_steps and global_step >= config.max_steps: break # Do not overwrite the source checkpoint if a resume exits without a new optimizer step. if not config.resume_from_checkpoint or global_step > _resume_initial_global_step: save_checkpoint( model, optimizer, scheduler, global_step, config, rank, epoch=epoch, micro_step=step + 1, mask_rng=mask_rng, loader_generator=loader_generator, sampler_seed=sampler_seed, ) elif is_main_process(rank): print(f"Skipping final save: no new optimizer step after resume at {global_step}") if config.wandb_enabled and is_main_process(rank): wandb.finish() if world_size > 1: dist.destroy_process_group() if is_main_process(rank): print("Training complete!") def main(): parser = argparse.ArgumentParser(description="Train Qwen3 model") parser.add_argument( "--config", type=str, default="/mloscratch/homes/navasard/sparse/SparseAtt/configs/train_config_dense_resume.json", help="Path to training config JSON", ) args = parser.parse_args() config = TrainConfig.from_json(args.config) train(config) if __name__ == "__main__": main()