| """ |
| 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) |
|
|
|
|
| |
| |
| |
|
|
| 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) |
|
|
| |
| 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.") |
|
|
| |
| total_w = sum(weights) |
| weights = [w / total_w for w in weights] |
|
|
| total_tokens = 0 |
|
|
| while any(not s["exhausted"] for s in streams): |
| |
| 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] |
|
|
| |
| 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 |
|
|
| |
| 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: |
| |
| if "instruction" in item and "output" in item: |
| |
| prompt = item["instruction"] |
| if item.get("input"): |
| prompt = f"{prompt}\n{item['input']}" |
| response = item["output"] |
| elif "messages" in item: |
| |
| messages = item["messages"] |
| if len(messages) >= 2: |
| prompt = messages[0].get("content", "") |
| response = messages[1].get("content", "") |
| else: |
| continue |
| elif "conversations" in item: |
| |
| 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 |
|
|
| |
| prompt_ids = tokenizer.encode(prompt, add_special_tokens=False) |
| response_ids = tokenizer.encode(response, add_special_tokens=False) |
|
|
| |
| 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) |
|
|
| |
| 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 |
|
|
|
|
| |
| |
| |
|
|
| 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.""" |
|
|
| |
| |
| 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) |
|
|
| |
| 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: |
| |
| 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) |
|
|
| |
| 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")) |
|
|
| |
| |
| |
| 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: |
| |
| 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() |
|
|
| |
| 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 |
|
|
|
|
| |
| |
| |
|
|
| 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}") |
|
|
| |
| model = PCSHODLM(config) |
| msa_layers = create_msa_layers(config, msa_config) |
|
|
| |
| 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):,}") |
|
|
| |
| 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") |
|
|
| |
| train_mode = args.mode |
| loader_workers = 2 if args.device.startswith("cuda") else 0 |
|
|
| if args.phase == "pretrain": |
| eval_batches = [] |
| if args.data_mix: |
| |
| |
| 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: |
| |
| 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() |
|
|
