Text Generation
Transformers
Safetensors
OpenPeerLLM
English
Arabic
riemann_openpeer
amd
ntk
code
fsdp
rocm
coding
mi325x
decoder
code gen
directml
openpeer
ryzen-ai
reasoning
rewriting
multimodal
paraphrase
prerelease
ntk-trainer
growth-ladder
style-transfer
text-rewriting
riemann-computing
progressive-growth
custom_code
Eval Results (legacy)
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import logging | |
| import sys | |
| import time | |
| from concurrent.futures import ThreadPoolExecutor | |
| from pathlib import Path | |
| sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) | |
| import torch | |
| from torch.utils.data import DataLoader | |
| from transformers import AutoTokenizer | |
| from transformers.optimization import Adafactor | |
| from riemann.dataset import PackedTextDataset | |
| from riemann.growth_ladder import get_stage, GrowthStage | |
| from riemann.growth import build_model | |
| from riemann.modeling_riemann import RiemannForCausalLM | |
| from riemann.openpeer_weights import build_stage0_from_openpeer_pretrained | |
| from riemann.device import ( | |
| main_device, | |
| secondary_device, | |
| cpu_device, | |
| configure_cpu_threads, | |
| describe_devices, | |
| ) | |
| logger = logging.getLogger("riemann.pretrain") | |
| from riemann.paths import DATA_DIR, CHECKPOINT_ROOT | |
| TOKENIZER_NAME = "gpt2" | |
| def load_or_build_model( | |
| stage: GrowthStage, checkpoint_dir: Path, resume: bool, use_pretrained_init: bool | |
| ): | |
| if resume and (checkpoint_dir / "config.json").exists(): | |
| logger.info("Resuming model weights from %s", checkpoint_dir) | |
| model = RiemannForCausalLM.from_pretrained(str(checkpoint_dir)) | |
| elif use_pretrained_init and stage.name == "stage0_124m": | |
| logger.info("Building %s from OpenPeerLLM's own published weights", stage.name) | |
| model = build_stage0_from_openpeer_pretrained(stage.to_riemann_config()) | |
| else: | |
| logger.info( | |
| "Building a fresh %s model (no checkpoint found / --resume not given)", | |
| stage.name, | |
| ) | |
| model = build_model(stage) | |
| return model | |
| def load_training_state(checkpoint_dir: Path) -> dict: | |
| state_path = checkpoint_dir / "training_state.json" | |
| if state_path.exists(): | |
| return json.loads(state_path.read_text(encoding="utf-8")) | |
| return {"stage": None, "global_step": 0} | |
| def save_checkpoint( | |
| model, tokenizer, checkpoint_dir: Path, stage: GrowthStage, global_step: int | |
| ) -> None: | |
| original_device = next(model.parameters()).device | |
| model.to("cpu") | |
| checkpoint_dir.mkdir(parents=True, exist_ok=True) | |
| model.save_pretrained(str(checkpoint_dir)) | |
| tokenizer.save_pretrained(str(checkpoint_dir)) | |
| model.to(original_device) | |
| state = {"stage": stage.name, "global_step": global_step} | |
| (checkpoint_dir / "training_state.json").write_text( | |
| json.dumps(state, indent=2), encoding="utf-8" | |
| ) | |
| logger.info("Saved checkpoint at step %d -> %s", global_step, checkpoint_dir) | |
| def run_cpu_microstep(cpu_model, batch, grad_accum: int) -> float: | |
| outputs = cpu_model(input_ids=batch, labels=batch) | |
| loss = outputs.loss / grad_accum | |
| loss.backward() | |
| return outputs.loss.item() | |
| def sync_cpu_mirror(cpu_model, model, cpu_dev) -> None: | |
| cpu_model.load_state_dict( | |
| {k: v.detach().to(cpu_dev) for k, v in model.state_dict().items()} | |
| ) | |
| for p in cpu_model.parameters(): | |
| p.grad = None | |
| def build_multimodal_step(model, adapter, device, secondary, tokenizer, example): | |
| frames = example["frames"].to(secondary) | |
| per_frame = adapter.frame_encoder(frames) | |
| pooled = per_frame.mean(dim=0, keepdim=True).to(device) | |
| video_tok = adapter.to_lm(pooled) | |
| prompt_ids = tokenizer.encode(example["prompt"])[:256] | |
| prompt_ids_t = torch.tensor(prompt_ids, dtype=torch.long, device=device) | |
| text_embeds = model.word_embeddings(prompt_ids_t).unsqueeze(0) | |
| inputs_embeds = torch.cat([video_tok.unsqueeze(0), text_embeds], dim=1) | |
| labels = torch.cat( | |
| [torch.tensor([-100], dtype=torch.long, device=device), prompt_ids_t] | |
| ).unsqueeze(0) | |
| outputs = model(inputs_embeds=inputs_embeds, labels=labels) | |
| return outputs.loss | |
| def main() -> None: | |
| logging.basicConfig( | |
| level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s" | |
| ) | |
| parser = argparse.ArgumentParser( | |
| description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter | |
| ) | |
| parser.add_argument( | |
| "--stage", default="stage0_124m", help="growth ladder stage name" | |
| ) | |
| parser.add_argument( | |
| "--resume", | |
| action="store_true", | |
| help="continue from this stage's saved checkpoint", | |
| ) | |
| parser.add_argument( | |
| "--max-steps", | |
| type=int, | |
| default=None, | |
| help="stop after this many optimizer steps (for smoke tests)", | |
| ) | |
| parser.add_argument( | |
| "--batch-size", | |
| type=int, | |
| default=None, | |
| help="override the stage's suggested batch size", | |
| ) | |
| parser.add_argument( | |
| "--grad-accum", | |
| type=int, | |
| default=None, | |
| help="override the stage's suggested grad accumulation", | |
| ) | |
| parser.add_argument( | |
| "--seq-len", | |
| type=int, | |
| default=None, | |
| help="override the stage's suggested sequence length", | |
| ) | |
| parser.add_argument( | |
| "--num-workers", | |
| type=int, | |
| default=2, | |
| help="DataLoader worker processes for tokenization", | |
| ) | |
| parser.add_argument("--log-every", type=int, default=10) | |
| parser.add_argument("--save-every", type=int, default=500) | |
| parser.add_argument( | |
| "--multimodal", | |
| action="store_true", | |
| help="also mix in CodeVideoBench video-conditioned steps", | |
| ) | |
| parser.add_argument( | |
| "--multimodal-every", | |
| type=int, | |
| default=50, | |
| help="run one multimodal step every N optimizer steps", | |
| ) | |
| parser.add_argument( | |
| "--no-pretrained-init", | |
| action="store_true", | |
| help="build stage0 from random weights instead of OpenPeerLLM's published checkpoint", | |
| ) | |
| parser.add_argument( | |
| "--no-cpu-accel", | |
| action="store_true", | |
| help="disable routing a share of micro-batches to the CPU during training", | |
| ) | |
| parser.add_argument( | |
| "--cpu-micro-step-every", | |
| type=int, | |
| default=4, | |
| help="1 out of every N micro-batches in a grad-accum cycle trains on CPU instead of the main GPU", | |
| ) | |
| parser.add_argument( | |
| "--cpu-threads", | |
| type=int, | |
| default=None, | |
| help="override CPU thread count (defaults to all logical cores)", | |
| ) | |
| parser.add_argument("--data-dir", default=str(DATA_DIR)) | |
| parser.add_argument("--checkpoint-root", default=str(CHECKPOINT_ROOT)) | |
| args = parser.parse_args() | |
| stage = get_stage(args.stage) | |
| batch_size = args.batch_size or stage.suggested_batch_size | |
| grad_accum = args.grad_accum or stage.suggested_grad_accum | |
| seq_len = args.seq_len or stage.suggested_seq_len | |
| checkpoint_dir = Path(args.checkpoint_root) / stage.name | |
| logger.info("Devices: %s", json.dumps(describe_devices())) | |
| if not stage.feasible_on_this_machine: | |
| logger.warning( | |
| "%s is marked NOT feasible on this machine. Proceeding anyway since you asked for it, but expect this to be slow or to run out of memory.", | |
| stage.name, | |
| ) | |
| device = main_device() | |
| tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME) | |
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| model = load_or_build_model( | |
| stage, checkpoint_dir, args.resume, not args.no_pretrained_init | |
| ) | |
| model.to(device) | |
| if stage.use_gradient_checkpointing: | |
| model.gradient_checkpointing_enable() | |
| model.config.use_cache = False | |
| model.train() | |
| cpu_dev = cpu_device() | |
| cpu_accel = (not args.no_cpu_accel) and device != cpu_dev | |
| cpu_model = None | |
| cpu_executor = None | |
| if cpu_accel: | |
| configure_cpu_threads(args.cpu_threads) | |
| cpu_model = type(model)(model.config) | |
| cpu_model.load_state_dict( | |
| {k: v.detach().to(cpu_dev) for k, v in model.state_dict().items()} | |
| ) | |
| cpu_model.to(cpu_dev) | |
| cpu_model.train() | |
| cpu_executor = ThreadPoolExecutor(max_workers=1) | |
| logger.info( | |
| "CPU acceleration enabled: 1 of every %d micro-batches trains on %s", | |
| args.cpu_micro_step_every, | |
| cpu_dev, | |
| ) | |
| state = load_training_state(checkpoint_dir) if args.resume else {"global_step": 0} | |
| global_step = state.get("global_step", 0) | |
| params = list(model.parameters()) | |
| adapter = None | |
| multimodal_examples = [] | |
| secondary = None | |
| if args.multimodal: | |
| from riemann.multimodal import ( | |
| VideoPrefixAdapter, | |
| build_multimodal_training_examples, | |
| ) | |
| secondary = secondary_device() | |
| adapter = VideoPrefixAdapter(lm_n_embd=model.config.hidden_size) | |
| adapter.frame_encoder.to(secondary) | |
| adapter.to_lm.to(device) | |
| params += list(adapter.parameters()) | |
| multimodal_examples = build_multimodal_training_examples() | |
| logger.info( | |
| "Multimodal enabled: %d CodeVideoBench examples, frame encoder on %s", | |
| len(multimodal_examples), | |
| secondary, | |
| ) | |
| optimizer = Adafactor( | |
| params, scale_parameter=True, relative_step=True, warmup_init=True, lr=None | |
| ) | |
| corpus_path = Path(args.data_dir) / "pretrain_corpus.txt" | |
| if not corpus_path.exists(): | |
| raise FileNotFoundError( | |
| f"{corpus_path} not found -- run scripts/prepare_datasets.py first" | |
| ) | |
| dataset = PackedTextDataset(corpus_path, tokenizer, seq_len) | |
| loader = DataLoader(dataset, batch_size=batch_size, num_workers=args.num_workers) | |
| logger.info( | |
| "Training %s: batch_size=%d grad_accum=%d seq_len=%d device=%s (resuming from step %d)", | |
| stage.name, | |
| batch_size, | |
| grad_accum, | |
| seq_len, | |
| device, | |
| global_step, | |
| ) | |
| mm_cursor = 0 | |
| accum_loss = 0.0 | |
| cpu_micro_steps = 0 | |
| micro_step = 0 | |
| pending_cpu_future = None | |
| start_time = time.time() | |
| optimizer.zero_grad() | |
| for batch in loader: | |
| micro_step += 1 | |
| route_to_cpu = cpu_accel and micro_step % args.cpu_micro_step_every == 0 | |
| if route_to_cpu: | |
| cpu_batch = batch.to(cpu_dev) | |
| if pending_cpu_future is not None: | |
| accum_loss += pending_cpu_future.result() | |
| cpu_micro_steps += 1 | |
| pending_cpu_future = cpu_executor.submit( | |
| run_cpu_microstep, cpu_model, cpu_batch, grad_accum | |
| ) | |
| else: | |
| batch = batch.to(device) | |
| outputs = model(input_ids=batch, labels=batch) | |
| loss = outputs.loss / grad_accum | |
| loss.backward() | |
| accum_loss += outputs.loss.item() | |
| if micro_step % grad_accum == 0: | |
| if pending_cpu_future is not None: | |
| accum_loss += pending_cpu_future.result() | |
| cpu_micro_steps += 1 | |
| pending_cpu_future = None | |
| for p, cp in zip(model.parameters(), cpu_model.parameters()): | |
| if cp.grad is None: | |
| continue | |
| cpu_grad = cp.grad.detach().to(device) | |
| p.grad = cpu_grad if p.grad is None else p.grad + cpu_grad | |
| optimizer.step() | |
| optimizer.zero_grad() | |
| if cpu_accel: | |
| sync_cpu_mirror(cpu_model, model, cpu_dev) | |
| global_step += 1 | |
| if ( | |
| adapter is not None | |
| and multimodal_examples | |
| and (global_step % args.multimodal_every == 0) | |
| ): | |
| example = multimodal_examples[mm_cursor % len(multimodal_examples)] | |
| mm_cursor += 1 | |
| mm_loss = build_multimodal_step( | |
| model, adapter, device, secondary, tokenizer, example | |
| ) | |
| (mm_loss / grad_accum).backward() | |
| optimizer.step() | |
| optimizer.zero_grad() | |
| logger.info( | |
| " [multimodal step on %s] loss=%.4f", | |
| example["name"], | |
| mm_loss.item(), | |
| ) | |
| if global_step % args.log_every == 0: | |
| elapsed = time.time() - start_time | |
| avg_loss = accum_loss / (args.log_every * grad_accum) | |
| logger.info( | |
| "step %d | loss %.4f | %.1fs elapsed | %d cpu micro-steps", | |
| global_step, | |
| avg_loss, | |
| elapsed, | |
| cpu_micro_steps, | |
| ) | |
| accum_loss = 0.0 | |
| cpu_micro_steps = 0 | |
| if global_step % args.save_every == 0: | |
| save_checkpoint(model, tokenizer, checkpoint_dir, stage, global_step) | |
| if args.max_steps is not None and global_step >= args.max_steps: | |
| break | |
| if cpu_executor is not None: | |
| cpu_executor.shutdown(wait=True) | |
| save_checkpoint(model, tokenizer, checkpoint_dir, stage, global_step) | |
| logger.info("Done. Final step %d, checkpoint at %s", global_step, checkpoint_dir) | |
| if __name__ == "__main__": | |
| main() | |