""" PC-SHO-DLM 2B Training Script — RunPod B200 (192GB VRAM) Three-phase training: Phase 1: Pretraining on FineWeb-Edu (50B tokens, BPE) Phase 2: Instruction tuning on UltraChat/Alpaca/Dolly (conditional masking) Phase 3: MSA router training on MS MARCO (contrastive routing) Usage: python train_2b.py --phase pretrain --steps 50000 python train_2b.py --phase instruct --steps 30000 --checkpoint results/pretrain_final.pt python train_2b.py --phase msa --steps 10000 --checkpoint results/instruct_final.pt """ import argparse, itertools, json, math, os, random, sys, time import shutil from pathlib import Path import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader, IterableDataset sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src")) from model import PCSHODLM, PCSHOConfig, LocalParameterUpdater, count_parameters from msa import MSAConfig, create_msa_layers, MemoryBank, MemoryEncoder, compute_routing_aux_loss MODEL_PRESETS = { "smoke": { "description": "Tiny MPS/CPU smoke-test preset", "config": dict( vocab_size=49153, max_seq_len=128, d_model=256, n_heads=8, n_layers=4, d_ff=1024, n_diffusion_steps=16, n_settling_steps=2, mask_token_id=49152, dropout=0.1, feedback_rank=64, denoising_lambda=5.0, homeostatic_target=0.0, lateral_inhibition=0.0, mass_scale=0.5, gamma_scale=0.3, eta_scale=0.3, ), "msa": dict( chunk_size=64, top_k=8, router_dim=64, n_router_heads=4, apply_from_layer=2, ), }, "dev": { "description": "Local/dev preset (~430M base params)", "config": dict( vocab_size=49153, max_seq_len=1024, d_model=1024, n_heads=8, n_layers=12, d_ff=4096, n_diffusion_steps=64, n_settling_steps=4, mask_token_id=49152, dropout=0.1, feedback_rank=256, denoising_lambda=5.0, homeostatic_target=0.0, lateral_inhibition=0.0, mass_scale=0.5, gamma_scale=0.3, eta_scale=0.3, ), "msa": dict( chunk_size=64, top_k=16, router_dim=128, n_router_heads=8, apply_from_layer=6, ), }, "2b": { "description": "RunPod B200 preset (~1.99B total params with MSA)", "config": dict( vocab_size=49153, max_seq_len=4096, d_model=1536, n_heads=12, n_layers=24, d_ff=6144, n_diffusion_steps=256, n_settling_steps=8, mask_token_id=49152, dropout=0.1, feedback_rank=384, denoising_lambda=5.0, homeostatic_target=0.0, lateral_inhibition=0.0, mass_scale=0.5, gamma_scale=0.3, eta_scale=0.3, ), "msa": dict( chunk_size=64, top_k=16, router_dim=128, n_router_heads=8, apply_from_layer=12, ), }, } _LOCAL_SNAPSHOT_CACHE = {} _LOCAL_DATA_SUFFIXES = {".parquet", ".jsonl", ".json"} def _validate_tokenizer_compatibility(config, tokenizer_name: str) -> None: from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) base_vocab_size = tokenizer.vocab_size supported_base_vocab = config.vocab_size - 1 if base_vocab_size > supported_base_vocab: raise SystemExit( f"Tokenizer '{tokenizer_name}' has vocab_size={base_vocab_size}, " f"but this preset only supports up to {supported_base_vocab} base tokens " f"(config.vocab_size={config.vocab_size}, mask_token_id={config.mask_token_id}). " "Use a compatible tokenizer such as 'HuggingFaceTB/SmolLM2-1.7B' or increase the model vocab." ) def _resolve_local_dataset_path(dataset_name): if os.path.isfile(dataset_name) or os.path.isdir(dataset_name): return dataset_name if dataset_name in _LOCAL_SNAPSHOT_CACHE: return _LOCAL_SNAPSHOT_CACHE[dataset_name] try: from huggingface_hub import snapshot_download path = snapshot_download( repo_id=dataset_name, repo_type="dataset", local_files_only=True, ) except Exception: path = None _LOCAL_SNAPSHOT_CACHE[dataset_name] = path return path def _select_local_data_files(dataset_name, split=None): base_path = _resolve_local_dataset_path(dataset_name) if base_path is None: return [] path = Path(base_path) if path.is_file(): return [path] files = sorted( p for p in path.rglob("*") if p.is_file() and p.suffix.lower() in _LOCAL_DATA_SUFFIXES ) if split: split_key = split.lower() matched = [ p for p in files if split_key in p.name.lower() or split_key in str(p.parent).lower() ] if matched: files = matched return files def _iter_json_array_file(path): decoder = json.JSONDecoder() buffer = "" in_array = False eof = False with open(path, "r", encoding="utf-8") as f: while True: if not eof and len(buffer) < (1 << 20): chunk = f.read(1 << 20) if chunk: buffer += chunk else: eof = True if not in_array: stripped = buffer.lstrip() if not stripped: if eof: return buffer = "" continue if stripped[0] == "[": in_array = True buffer = stripped[1:] else: payload = json.loads(stripped) if isinstance(payload, list): for item in payload: yield item else: yield payload return buffer = buffer.lstrip() if not buffer: if eof: return continue if buffer[0] == "]": return if buffer[0] == ",": buffer = buffer[1:] continue try: item, next_idx = decoder.raw_decode(buffer) except json.JSONDecodeError: if eof: raise continue yield item buffer = buffer[next_idx:] def _iter_local_records(dataset_name, split=None): files = _select_local_data_files(dataset_name, split=split) if not files: return None def iterator(): for path in files: suffix = path.suffix.lower() if suffix == ".parquet": import pyarrow.parquet as pq parquet_file = pq.ParquetFile(path) for batch in parquet_file.iter_batches(batch_size=64, use_threads=False): for row in batch.to_pylist(): yield row elif suffix == ".jsonl": with open(path, "r", encoding="utf-8") as f: for line in f: line = line.strip() if line: yield json.loads(line) elif suffix == ".json": yield from _iter_json_array_file(path) return iterator() def _open_record_stream(dataset_name, config_name=None, split="train"): local_records = _iter_local_records(dataset_name, split=split) if local_records is not None: return local_records from datasets import load_dataset ds = load_dataset(dataset_name, config_name, split=split, streaming=True) return iter(ds) def _shuffle_stream(stream, buffer_size: int, seed: int | None = None): """Approximate shuffle for streaming iterables using a bounded buffer.""" if buffer_size <= 1: yield from stream return rng = random.Random(seed) buffer = [] for item in stream: if len(buffer) < buffer_size: buffer.append(item) continue idx = rng.randrange(len(buffer)) yield buffer[idx] buffer[idx] = item rng.shuffle(buffer) for item in buffer: yield item def _shard_stream(stream, worker_id: int, num_workers: int): """Shard an iterable stream across dataloader workers.""" if num_workers <= 1: yield from stream return yield from itertools.islice(stream, worker_id, None, num_workers) # ============================================================================= # BPE Tokenized Streaming Dataset # ============================================================================= class BPEStreamDataset(IterableDataset): """Streams BPE-tokenized text from HuggingFace datasets.""" def __init__(self, dataset_name: str, config_name: str, split: str, tokenizer_name: str, max_seq_len: int, max_tokens: int = None, mask_token_id: int = None, text_field: str = "text", shuffle_buffer_size: int = 0, shuffle_seed: int = 0, record_skip: int = 0, max_sequences: int | None = None): self.dataset_name = dataset_name self.config_name = config_name self.split = split self.tokenizer_name = tokenizer_name self.max_seq_len = max_seq_len self.max_tokens = max_tokens self.mask_token_id = mask_token_id self.text_field = text_field self.shuffle_buffer_size = shuffle_buffer_size self.shuffle_seed = shuffle_seed self.record_skip = record_skip self.max_sequences = max_sequences def __iter__(self): from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.model_max_length = max(tokenizer.model_max_length, 10**30) if os.path.isfile(self.dataset_name): yield from self._iter_local_text_file(tokenizer) return buffer = [] total = 0 worker_info = torch.utils.data.get_worker_info() worker_offset = worker_info.id if worker_info else 0 worker_count = worker_info.num_workers if worker_info else 1 record_stream = _open_record_stream(self.dataset_name, self.config_name, self.split) record_stream = _shard_stream(record_stream, worker_offset, worker_count) if self.record_skip > 0: record_stream = itertools.islice(record_stream, self.record_skip, None) if self.shuffle_buffer_size > 1: record_stream = _shuffle_stream( record_stream, self.shuffle_buffer_size, seed=self.shuffle_seed + worker_offset, ) yielded = 0 for item in record_stream: text = item.get(self.text_field, "") if not text or not text.strip(): continue tokens = tokenizer.encode(text, add_special_tokens=False) buffer.extend(tokens) while len(buffer) >= self.max_seq_len: chunk = buffer[:self.max_seq_len] buffer = buffer[self.max_seq_len:] total += self.max_seq_len if self.max_tokens and total > self.max_tokens: return yield {"input_ids": torch.tensor(chunk, dtype=torch.long)} yielded += 1 if self.max_sequences is not None and yielded >= self.max_sequences: return def _iter_local_text_file(self, tokenizer): buffer = [] total = 0 with open(self.dataset_name, "r", encoding="utf-8", errors="ignore") as f: for line in f: text = line.strip() if not text: continue tokens = tokenizer.encode(text, add_special_tokens=False) buffer.extend(tokens) while len(buffer) >= self.max_seq_len: chunk = buffer[:self.max_seq_len] buffer = buffer[self.max_seq_len:] total += self.max_seq_len if self.max_tokens and total > self.max_tokens: return yield {"input_ids": torch.tensor(chunk, dtype=torch.long)} class InterleavedStreamDataset(IterableDataset): """Interleaves multiple BPEStreamDatasets with weighted sampling. Usage: dataset = InterleavedStreamDataset( sources=[ {"name": "HuggingFaceFW/fineweb-edu", "config": "sample-10BT", "weight": 0.50}, {"name": "HuggingFaceTB/cosmopedia-v2", "config": "cosmopedia-v2", "weight": 0.20}, {"name": "open-web-math/open-web-math", "weight": 0.15}, {"name": "bigcode/self-oss-instruct-sc2-exec-filter-50k", "text_field": "response", "weight": 0.10}, {"name": "HuggingFaceTB/finemath", "config": "finemath-3plus", "weight": 0.05}, ], tokenizer_name="gpt2", max_seq_len=1024, max_tokens=50_000_000_000, ) """ def __init__(self, sources: list, tokenizer_name: str, max_seq_len: int, max_tokens: int = None, mask_token_id: int = None, shuffle_buffer_size: int = 0, shuffle_seed: int = 0): self.sources = sources self.tokenizer_name = tokenizer_name self.max_seq_len = max_seq_len self.max_tokens = max_tokens self.mask_token_id = mask_token_id self.shuffle_buffer_size = shuffle_buffer_size self.shuffle_seed = shuffle_seed def __iter__(self): from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.model_max_length = max(tokenizer.model_max_length, 10**30) worker_info = torch.utils.data.get_worker_info() worker_offset = worker_info.id if worker_info else 0 worker_count = worker_info.num_workers if worker_info else 1 rng = random.Random(self.shuffle_seed + worker_offset) # Open all streams streams = [] weights = [] for src_idx, src in enumerate(self.sources): try: ds_iter = _open_record_stream( src["name"], src.get("config"), split=src.get("split", "train"), ) ds_iter = _shard_stream(ds_iter, worker_offset, worker_count) if self.shuffle_buffer_size > 1: ds_iter = _shuffle_stream( ds_iter, self.shuffle_buffer_size, seed=self.shuffle_seed + worker_offset + src_idx, ) except Exception as e: print(f"[WARN] Skipping data source {src['name']}: {e}") continue streams.append({ "iter": iter(ds_iter), "text_field": src.get("text_field", "text"), "buffer": [], "exhausted": False, "name": src["name"], }) weights.append(src.get("weight", 1.0)) if not streams: raise RuntimeError("No pretraining sources are available locally or remotely.") # Normalize weights total_w = sum(weights) weights = [w / total_w for w in weights] total_tokens = 0 while any(not s["exhausted"] for s in streams): # Weighted random source selection (skip exhausted) active = [(i, w) for i, (s, w) in enumerate(zip(streams, weights)) if not s["exhausted"]] if not active: break indices, ws = zip(*active) w_sum = sum(ws) probs = [w / w_sum for w in ws] idx = rng.choices(indices, weights=probs, k=1)[0] stream = streams[idx] # Fill buffer from selected stream while len(stream["buffer"]) < self.max_seq_len: try: item = next(stream["iter"]) text = item.get(stream["text_field"], "") if text and text.strip(): tokens = tokenizer.encode(text, add_special_tokens=False) stream["buffer"].extend(tokens) except StopIteration: stream["exhausted"] = True break # Yield chunks from buffer while len(stream["buffer"]) >= self.max_seq_len: chunk = stream["buffer"][:self.max_seq_len] stream["buffer"] = stream["buffer"][self.max_seq_len:] total_tokens += self.max_seq_len if self.max_tokens and total_tokens > self.max_tokens: return yield {"input_ids": torch.tensor(chunk, dtype=torch.long)} class DialogueStreamDataset(IterableDataset): """Streams instruction-response pairs with conditional masking.""" def __init__(self, dataset_name: str, split: str, tokenizer_name: str, max_seq_len: int, max_samples: int = None): self.dataset_name = dataset_name self.split = split self.tokenizer_name = tokenizer_name self.max_seq_len = max_seq_len self.max_samples = max_samples def __iter__(self): from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.model_max_length = max(tokenizer.model_max_length, 10**30) count = 0 worker_info = torch.utils.data.get_worker_info() worker_offset = worker_info.id if worker_info else 0 worker_count = worker_info.num_workers if worker_info else 1 record_stream = _open_record_stream(self.dataset_name, split=self.split) record_stream = _shard_stream(record_stream, worker_offset, worker_count) for item in record_stream: # Handle various dataset formats if "instruction" in item and "output" in item: # Alpaca-style: instruction + optional input → output prompt = item["instruction"] if item.get("input"): prompt = f"{prompt}\n{item['input']}" response = item["output"] elif "messages" in item: # UltraChat-style: [{"role": ..., "content": ...}, ...] messages = item["messages"] if len(messages) >= 2: prompt = messages[0].get("content", "") response = messages[1].get("content", "") else: continue elif "conversations" in item: # OpenHermes/SlimOrca-style: [{"from": ..., "value": ...}, ...] convs = item["conversations"] if len(convs) >= 2: prompt = convs[0].get("value", "") response = convs[1].get("value", "") else: continue elif "prompt" in item and "response" in item: prompt = item["prompt"] response = item["response"] else: continue if not prompt or not response: continue # Tokenize prompt_ids = tokenizer.encode(prompt, add_special_tokens=False) response_ids = tokenizer.encode(response, add_special_tokens=False) # Truncate max_prompt = self.max_seq_len // 2 prompt_ids = prompt_ids[:max_prompt] response_ids = response_ids[:self.max_seq_len - len(prompt_ids)] full = prompt_ids + response_ids prompt_len = len(prompt_ids) # Pad if len(full) < self.max_seq_len: full = full + [tokenizer.pad_token_id] * (self.max_seq_len - len(full)) count += 1 if self.max_samples and count > self.max_samples: return yield { "input_ids": torch.tensor(full[:self.max_seq_len], dtype=torch.long), "prompt_len": prompt_len, } def _build_fixed_eval_batches( dataset_name: str, config_name: str | None, tokenizer_name: str, config, *, batch_size: int, n_batches: int, record_skip: int, shuffle_buffer_size: int, shuffle_seed: int, mask_seed: int = 1234, ): if n_batches <= 0: return [] dataset = BPEStreamDataset( dataset_name, config_name, "train", tokenizer_name, config.max_seq_len, mask_token_id=config.mask_token_id, shuffle_buffer_size=shuffle_buffer_size, shuffle_seed=shuffle_seed, record_skip=record_skip, max_sequences=n_batches * batch_size, ) loader = DataLoader(dataset, batch_size=batch_size, num_workers=0) rng = torch.Generator(device="cpu") rng.manual_seed(mask_seed) alpha_bar = torch.cos( (torch.arange(config.n_diffusion_steps + 1, dtype=torch.float32) / config.n_diffusion_steps) * (math.pi / 2) ) ** 2 eval_batches = [] for batch in loader: x_0 = batch["input_ids"].clone() B, S = x_0.shape t = torch.randint(1, config.n_diffusion_steps + 1, (B,), generator=rng) mask_prob = (1.0 - alpha_bar[t]).unsqueeze(-1) mask = torch.rand((B, S), generator=rng) < mask_prob empty_rows = ~mask.any(dim=1) if empty_rows.any(): forced = torch.randint(0, S, (int(empty_rows.sum().item()),), generator=rng) mask[empty_rows] = False mask[empty_rows, forced] = True x_t = x_0.clone() x_t[mask] = config.mask_token_id eval_batches.append({"input_ids": x_0, "x_t": x_t, "mask": mask, "t": t}) if len(eval_batches) >= n_batches: break return eval_batches def _load_matching_state_dict(model: nn.Module, state_dict: dict) -> tuple[list[str], list[str]]: """Load only checkpoint tensors whose shapes match the current model.""" current = model.state_dict() matched = {} skipped = [] missing = [] for key, value in state_dict.items(): if key not in current: skipped.append(key) continue if current[key].shape != value.shape: skipped.append(key) continue matched[key] = value for key in current.keys(): if key not in matched: missing.append(key) current.update(matched) model.load_state_dict(current, strict=False) return skipped, missing # ============================================================================= # Training Loop # ============================================================================= def train_phase(model, dataloader, mode, max_steps, lr, save_dir, phase_name, device="cuda", log_interval=100, save_interval=5000, msa_layers=None, msa_config=None, start_step=0, hf_upload=False, hf_repo_id=None, warmup_steps=None, min_lr_scale=0.1, tensorboard_dir=None, tensorboard_flush_secs=30, eval_batches=None, eval_interval=0): """Generic training loop for any phase. Uses DataParallel for multi-GPU.""" # Multi-GPU via Local SGD: split batch across GPUs, run unified settling # independently, average parameters after each step. import copy use_cuda = str(device).startswith("cuda") if use_cuda and not torch.cuda.is_available(): raise RuntimeError( "CUDA was requested but no CUDA devices are visible. " "Run on a GPU node or pass --device cpu for a small smoke test." ) n_gpu = torch.cuda.device_count() if use_cuda else 0 primary_device = "cuda:0" if use_cuda else device device = primary_device model = model.to(primary_device) raw_model = model raw_model.train() n_params = count_parameters(raw_model) # Create replica on GPU 1 for multi-GPU model_replica = None if use_cuda and n_gpu > 1 and mode == "unified": model_replica = copy.deepcopy(model).to("cuda:1") model_replica.train() print(f"Multi-GPU: Local SGD across {n_gpu} GPUs") if mode == "local": updater = LocalParameterUpdater( raw_model, lr_forward=lr, lr_feedback=lr, lr_readout=lr, lr_precision=lr * 0.1 ) elif mode == "backprop": optimizer = torch.optim.AdamW(raw_model.parameters(), lr=lr, weight_decay=0.01) os.makedirs(save_dir, exist_ok=True) writer = None if tensorboard_dir: try: from torch.utils.tensorboard import SummaryWriter os.makedirs(tensorboard_dir, exist_ok=True) writer = SummaryWriter( log_dir=tensorboard_dir, flush_secs=tensorboard_flush_secs, purge_step=start_step if start_step > 0 else None, ) except Exception as e: print(f"[WARN] TensorBoard disabled: {e}") step = start_step start = time.time() loss_ema = None auto_warmup_steps = min(2000, max(10, max_steps // 10)) if warmup_steps is None: warmup_steps = auto_warmup_steps reserve_disk_gb = max(5.0, (n_params * 4 / (1024 ** 3)) + 2.0) def lr_for_step(step_idx): if max_steps <= 1: return lr if warmup_steps > 0 and step_idx < warmup_steps: return lr * ((step_idx + 1) / warmup_steps) progress = (step_idx - warmup_steps) / max(1, max_steps - warmup_steps) progress = min(max(progress, 0.0), 1.0) cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) return lr * (min_lr_scale + (1.0 - min_lr_scale) * cosine) def evaluate_fixed_batches(): if not eval_batches: return None raw_model.eval() losses = [] with torch.no_grad(): for eval_batch in eval_batches: x_0_eval = eval_batch["input_ids"].to(device) x_t_eval = eval_batch["x_t"].to(device) mask_eval = eval_batch["mask"].to(device) t_eval = eval_batch["t"].to(device) h_0_eval = raw_model.embed_input(x_t_eval, t_eval) h_init_eval = raw_model.amortized_forward_pass(h_0_eval) eval_loss = raw_model._masked_cross_entropy( h_init_eval[raw_model.config.n_layers], x_0_eval, mask_eval ) losses.append(float(eval_loss.item())) raw_model.train() return sum(losses) / len(losses) if losses else None print(f"\n{'='*70}") print(f"{phase_name} | {n_params:,} params | {mode} mode | {device} | {n_gpu} GPUs") if start_step > 0: print(f"Resuming from step {start_step}") print(f"LR schedule: base={lr:.2e} warmup_steps={warmup_steps} min_scale={min_lr_scale:.2f}") print(f"{'='*70}\n") try: for batch in dataloader: if step >= max_steps: break x_0 = batch["input_ids"].to(device) B, S = x_0.shape step_lr = lr_for_step(step) settled_loss_val = None if mode == "unified": if model_replica is not None and B >= 2: # === MULTI-GPU: Local SGD === half = B // 2 x_0_a, x_0_b = x_0[:half].to("cuda:0"), x_0[half:].to("cuda:1") result_a = raw_model.unified_train_batch(x_0_a, param_lr=step_lr) energies = result_a.get("energies", []) result_b = model_replica.unified_train_batch(x_0_b, param_lr=step_lr) # Average parameters across GPUs with torch.no_grad(): for p0, p1 in zip(raw_model.parameters(), model_replica.parameters()): avg = (p0.data + p1.data.to("cuda:0")) / 2 p0.data.copy_(avg) p1.data.copy_(avg.to("cuda:1")) # Unified training returns a stale settled loss and a fresh # amortized loss. The amortized loss is the cleaner progress # metric for current parameters. loss_val = result_a.get("amortized_loss", result_a.get("loss", 0.0)) settled_loss_val = result_a.get("settled_loss", result_a.get("loss")) else: # === SINGLE GPU === result = raw_model.unified_train_batch(x_0, param_lr=step_lr) loss_val = result.get("amortized_loss", result.get("loss", 0.0)) settled_loss_val = result.get("settled_loss", result.get("loss")) energies = result.get("energies", []) if math.isnan(loss_val) or math.isinf(loss_val): print(f" [WARN] NaN/Inf at step {step}") loss_val = 0.0 elif mode == "local": updater.set_learning_rates( lr_forward=step_lr, lr_feedback=step_lr, lr_readout=step_lr, lr_precision=step_lr * 0.1, lr_embed=step_lr, ) if "prompt_len" in batch: prompt_lens = batch["prompt_len"] t = torch.randint(1, raw_model.config.n_diffusion_steps + 1, (B,), device=device) mask_prob = raw_model.schedule.mask_rate(t).unsqueeze(-1) random_mask = torch.rand_like(x_0.float()) < mask_prob response_mask = torch.zeros_like(x_0, dtype=torch.bool) for b in range(B): response_mask[b, prompt_lens[b]:] = True mask = random_mask & response_mask x_t = x_0.clone() x_t[mask] = raw_model.config.mask_token_id h_0 = raw_model.embed_input(x_t, t) h_init = raw_model.amortized_forward_pass(h_0) with torch.no_grad(): h_settled, _, energies, _, _ = raw_model.settle(h_init, x_0, mask, t) updater.local_update(h_settled, x_0, mask, t) with torch.no_grad(): logits = raw_model.readout(raw_model.readout_norm(h_settled[raw_model.config.n_layers])) ml, mt = logits[mask], x_0[mask] loss_val = F.cross_entropy(ml, mt).item() if ml.numel() > 0 else 0.0 else: result = updater.step({"input_ids": x_0}) loss_val = result.get("loss", 0.0) energies = result.get("energies", []) elif mode == "backprop": optimizer.zero_grad() for group in optimizer.param_groups: group["lr"] = step_lr output = raw_model(x_0) loss = output["loss"] if torch.isnan(loss) or torch.isinf(loss): print(f" [WARN] NaN/Inf loss at step {step}, skipping") step += 1 continue loss.backward() torch.nn.utils.clip_grad_norm_(raw_model.parameters(), 1.0) optimizer.step() loss_val = loss.item() energies = output.get("energies", []) step += 1 loss_ema = loss_val if loss_ema is None else 0.95 * loss_ema + 0.05 * loss_val elapsed = time.time() - start steps_done = step - start_step tps = steps_done * B * S / elapsed if elapsed > 0 else 0 if writer is not None: writer.add_scalar(f"{phase_name}/loss", loss_val, step) writer.add_scalar(f"{phase_name}/loss_ema", loss_ema, step) writer.add_scalar(f"{phase_name}/lr", step_lr, step) writer.add_scalar(f"{phase_name}/tokens_per_sec", tps, step) if energies: writer.add_scalar(f"{phase_name}/energy", float(energies[-1]), step) if mode == "unified" and settled_loss_val is not None: writer.add_scalar(f"{phase_name}/settled_loss", settled_loss_val, step) if eval_interval and eval_batches and step % eval_interval == 0: eval_loss = evaluate_fixed_batches() if eval_loss is not None: print(f"[{phase_name}] Eval {step:6d} | HeldoutLoss: {eval_loss:.4f}") if writer is not None: writer.add_scalar(f"{phase_name}/heldout_loss", eval_loss, step) if step % log_interval == 0: e = f"{energies[-1]:.0f}" if energies else "N/A" metric_label = "Loss" metric_bits = [f"{metric_label}: {loss_val:.4f}", f"LossEMA: {loss_ema:.4f}"] if mode == "unified" and settled_loss_val is not None: metric_bits.append(f"Settled: {settled_loss_val:.4f}") metrics = " | ".join(metric_bits) print(f"[{phase_name}] Step {step:6d} | {metrics} | Energy: {e} | " f"LR: {step_lr:.2e} | Tok/s: {tps:.0f} | {elapsed:.0f}s") if save_interval and step % save_interval == 0: free_gb = shutil.disk_usage(save_dir).free / (1024 ** 3) if writer is not None: writer.add_scalar("system/free_disk_gb", free_gb, step) if free_gb < reserve_disk_gb: print( f" [WARN] Skipping checkpoint at step {step}: " f"only {free_gb:.1f} GiB free, need >= {reserve_disk_gb:.1f} GiB" ) continue path = os.path.join(save_dir, f"{phase_name}_step{step}.pt") torch.save({ "model_state_dict": raw_model.state_dict(), "config": raw_model.config, "step": step, "phase": phase_name, "loss": loss_val, }, path) print(f" Checkpoint saved: {path}") if hf_upload: try: from huggingface_hub import HfApi api = HfApi() repo_id = hf_repo_id or "zotowata/pc-sho-dlm-1.3b" api.create_repo(repo_id, exist_ok=True, private=True) api.upload_file( path_or_fileobj=path, path_in_repo=f"checkpoints/{phase_name}_step{step}.pt", repo_id=repo_id, ) print(f" Uploaded to HuggingFace: {repo_id}") except Exception as e: print(f" HF upload failed: {e}") finally: if writer is not None: writer.flush() writer.close() # Final save path = os.path.join(save_dir, f"{phase_name}_final.pt") free_gb = shutil.disk_usage(save_dir).free / (1024 ** 3) if free_gb < reserve_disk_gb: print( f"\n[WARN] Skipping final save: only {free_gb:.1f} GiB free, " f"need >= {reserve_disk_gb:.1f} GiB" ) path = None else: torch.save({ "model_state_dict": raw_model.state_dict(), "config": raw_model.config, "step": step, "phase": phase_name, }, path) if hf_upload: try: from huggingface_hub import HfApi api = HfApi() repo_id = hf_repo_id or "zotowata/pc-sho-dlm-1.3b" api.create_repo(repo_id, exist_ok=True, private=True) api.upload_file( path_or_fileobj=path, path_in_repo=f"checkpoints/{phase_name}_final.pt", repo_id=repo_id, ) print(f"Uploaded final checkpoint to Hugging Face: {repo_id}") except Exception as e: print(f"HF final upload failed: {e}") elapsed = time.time() - start print(f"\n{phase_name} complete: {step} steps in {elapsed:.0f}s") if path: print(f"Saved: {path}") return model # ============================================================================= # Main # ============================================================================= def main(): parser = argparse.ArgumentParser(description="PC-SHO-DLM 2B Training") parser.add_argument("--phase", choices=["pretrain", "instruct", "msa"], required=True) parser.add_argument("--steps", type=int, default=50000) parser.add_argument("--checkpoint", type=str, default=None) parser.add_argument("--lr", type=float, default=3e-4) parser.add_argument("--batch_size", type=int, default=8) parser.add_argument("--device", type=str, default="cuda") parser.add_argument("--save_dir", type=str, default="results") parser.add_argument("--tokenizer", type=str, default="HuggingFaceTB/SmolLM2-1.7B") parser.add_argument("--data", type=str, default=None) parser.add_argument("--model-preset", choices=sorted(MODEL_PRESETS.keys()), default="dev", help="Architecture preset to use") parser.add_argument("--log-interval", type=int, default=100, help="How often to print training metrics") parser.add_argument("--save-interval", type=int, default=5000, help="How often to write checkpoints during training") parser.add_argument("--warmup-steps", type=int, default=None, help="Linear warmup steps before cosine decay (default: auto)") parser.add_argument("--min-lr-scale", type=float, default=0.1, help="Minimum cosine-decayed LR as a fraction of base LR") parser.add_argument("--hf-upload", action="store_true", help="Upload periodic checkpoints to Hugging Face") parser.add_argument("--hf-repo-id", type=str, default=None, help="Override Hugging Face repo id for checkpoint uploads") parser.add_argument("--data-mix", action="store_true", help="Use interleaved multi-source pretrain mix instead of single dataset") parser.add_argument("--mode", choices=["unified", "backprop", "local"], default="unified", help="Training mode (default: unified)") parser.add_argument("--pretrain-shuffle-buffer", type=int, default=4096, help="Approximate record-shuffle buffer for pretraining streams") parser.add_argument("--pretrain-shuffle-seed", type=int, default=17, help="Random seed for pretraining stream shuffling") parser.add_argument("--seq-len-override", type=int, default=None, help="Optionally shorten max sequence length for curriculum training") parser.add_argument("--eval-interval", type=int, default=500, help="How often to run held-out evaluation during training") parser.add_argument("--eval-batches", type=int, default=8, help="Number of fixed held-out batches to evaluate") parser.add_argument("--eval-batch-size", type=int, default=1, help="Batch size for held-out evaluation") parser.add_argument("--eval-record-skip", type=int, default=200000, help="How many source records to skip before sampling held-out eval batches") parser.add_argument("--tensorboard-dir", type=str, default=None, help="Optional TensorBoard log directory") parser.add_argument("--tensorboard-flush-secs", type=int, default=30, help="TensorBoard event flush interval in seconds") args = parser.parse_args() if args.device.startswith("cuda") and not torch.cuda.is_available(): raise SystemExit( "No CUDA devices detected on this machine. " "A real 2B run needs a GPU node; use --device cpu only for tiny smoke tests." ) preset = MODEL_PRESETS[args.model_preset] config = PCSHOConfig(**preset["config"]) msa_config = MSAConfig(**preset["msa"]) if args.seq_len_override is not None: if args.seq_len_override <= 0 or args.seq_len_override > config.max_seq_len: raise SystemExit( f"--seq-len-override must be between 1 and {config.max_seq_len} for preset {args.model_preset}" ) config.max_seq_len = args.seq_len_override _validate_tokenizer_compatibility(config, args.tokenizer) save_root = args.save_dir if os.path.isabs(args.save_dir) else os.path.join(os.getcwd(), args.save_dir) tensorboard_root = None if args.tensorboard_dir: tensorboard_root = ( args.tensorboard_dir if os.path.isabs(args.tensorboard_dir) else os.path.join(os.getcwd(), args.tensorboard_dir) ) free_bytes = shutil.disk_usage(os.path.dirname(save_root) or os.getcwd()).free free_gb = free_bytes / (1024 ** 3) if free_gb < 10: print(f"[WARN] Low disk space: {free_gb:.1f} GiB free near save_dir={args.save_dir}") # Create model model = PCSHODLM(config) msa_layers = create_msa_layers(config, msa_config) # Load checkpoint if provided start_step = 0 if args.checkpoint: print(f"Loading checkpoint: {args.checkpoint}") ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False) skipped_keys, missing_keys = _load_matching_state_dict(model, ckpt["model_state_dict"]) start_step = ckpt.get("step", 0) print(f" Resuming from step {start_step}") if skipped_keys: print(f" Skipped {len(skipped_keys)} mismatched checkpoint tensors") if missing_keys: print(f" Current model has {len(missing_keys)} tensors initialized from scratch") total_params = count_parameters(model) + count_parameters(msa_layers) print(f"Preset: {args.model_preset} ({preset['description']})") print(f"Model: {total_params:,} total params") print(f" Base: {count_parameters(model):,}") print(f" MSA: {count_parameters(msa_layers):,}") # Enable bf16 for memory efficiency if args.device == "cuda": model = model.to(torch.bfloat16) print("Using bfloat16") print(f"Estimated memory: ~{count_parameters(model) * 12 / 1e9:.1f}GB for params+optim") # Phase-specific setup train_mode = args.mode loader_workers = 2 if args.device.startswith("cuda") else 0 if args.phase == "pretrain": eval_batches = [] if args.data_mix: # ── Multi-source interleaved pretraining mix ── # Weights follow the research mix: web 50%, synth-edu 20%, math 15%, code 10%, math-filtered 5% print("Using interleaved multi-source pretrain mix:") sources = [ {"name": "HuggingFaceFW/fineweb-edu", "config": "sample-10BT", "weight": 0.50, "text_field": "text"}, {"name": "HuggingFaceTB/cosmopedia-v2", "config": "cosmopedia-v2", "weight": 0.20, "text_field": "text"}, {"name": "open-web-math/open-web-math", "weight": 0.15, "text_field": "text"}, {"name": "bigcode/self-oss-instruct-sc2-exec-filter-50k", "weight": 0.10, "text_field": "response"}, {"name": "HuggingFaceTB/finemath", "config": "finemath-3plus", "weight": 0.05, "text_field": "text"}, ] for s in sources: print(f" {s['weight']*100:4.0f}% {s['name']}") dataset = InterleavedStreamDataset( sources=sources, tokenizer_name=args.tokenizer, max_seq_len=config.max_seq_len, max_tokens=50_000_000_000, mask_token_id=config.mask_token_id, shuffle_buffer_size=args.pretrain_shuffle_buffer, shuffle_seed=args.pretrain_shuffle_seed, ) else: # ── Single-source pretraining ── data_name = args.data or "HuggingFaceFW/fineweb-edu" data_config = "sample-10BT" if "fineweb" in data_name else None dataset = BPEStreamDataset( data_name, data_config, "train", args.tokenizer, config.max_seq_len, max_tokens=50_000_000_000, mask_token_id=config.mask_token_id, shuffle_buffer_size=args.pretrain_shuffle_buffer, shuffle_seed=args.pretrain_shuffle_seed, ) eval_batches = _build_fixed_eval_batches( data_name, data_config, args.tokenizer, config, batch_size=args.eval_batch_size, n_batches=args.eval_batches, record_skip=args.eval_record_skip, shuffle_buffer_size=args.pretrain_shuffle_buffer, shuffle_seed=args.pretrain_shuffle_seed + 1000, ) if eval_batches: print( f"Prepared {len(eval_batches)} held-out eval batches " f"(batch_size={args.eval_batch_size}, record_skip={args.eval_record_skip})" ) loader = DataLoader(dataset, batch_size=args.batch_size, num_workers=loader_workers) model = train_phase(model, loader, train_mode, args.steps, args.lr, args.save_dir, "pretrain", args.device, log_interval=args.log_interval, save_interval=args.save_interval, start_step=start_step, hf_upload=args.hf_upload, hf_repo_id=args.hf_repo_id, warmup_steps=args.warmup_steps, min_lr_scale=args.min_lr_scale, tensorboard_dir=tensorboard_root, tensorboard_flush_secs=args.tensorboard_flush_secs, eval_batches=eval_batches, eval_interval=args.eval_interval) elif args.phase == "instruct": data_name = args.data or "HuggingFaceH4/ultrachat_200k" dataset = DialogueStreamDataset( data_name, "train_sft", args.tokenizer, config.max_seq_len, max_samples=2_000_000, ) loader = DataLoader(dataset, batch_size=args.batch_size, num_workers=loader_workers) model = train_phase(model, loader, train_mode, args.steps, args.lr * 0.3, args.save_dir, "instruct", args.device, log_interval=args.log_interval, save_interval=args.save_interval, start_step=start_step, hf_upload=args.hf_upload, hf_repo_id=args.hf_repo_id, warmup_steps=args.warmup_steps, min_lr_scale=args.min_lr_scale, tensorboard_dir=tensorboard_root, tensorboard_flush_secs=args.tensorboard_flush_secs) elif args.phase == "msa": print("MSA Router Training — requires document QA dataset") print("TODO: Implement MSA-specific training loop with contrastive routing loss") print("Use MS MARCO or Natural Questions with document evidence labels") if __name__ == "__main__": main()