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 functools | |
| import json | |
| import logging | |
| import math | |
| import os | |
| import sys | |
| import time | |
| from pathlib import Path | |
| sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) | |
| import torch | |
| import torch.distributed as dist | |
| from torch.nn.parallel import DistributedDataParallel as DDP | |
| from torch.distributed.fsdp import ( | |
| FullyShardedDataParallel as FSDP, | |
| MixedPrecision, | |
| ShardingStrategy, | |
| FullStateDictConfig, | |
| StateDictType, | |
| ) | |
| from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy | |
| from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( | |
| checkpoint_wrapper, | |
| apply_activation_checkpointing, | |
| CheckpointImpl, | |
| ) | |
| from torch.utils.data import DataLoader | |
| from transformers import AutoTokenizer | |
| from riemann.dataset import PackedTextDataset | |
| from riemann.growth_ladder import get_stage, GrowthStage | |
| from riemann.modeling_riemann import RiemannForCausalLM, TransformerBlock | |
| from riemann.paths import DATA_DIR, CHECKPOINT_ROOT | |
| logger = logging.getLogger("riemann.pretrain_cuda") | |
| TOKENIZER_NAME = "gpt2" | |
| def setup_distributed(cpu: bool) -> tuple[int, int, int]: | |
| if cpu: | |
| dist.init_process_group(backend="gloo") | |
| return dist.get_rank(), dist.get_world_size(), 0 | |
| dist.init_process_group(backend="nccl") | |
| rank = dist.get_rank() | |
| world = dist.get_world_size() | |
| local_rank = int(os.environ.get("LOCAL_RANK", rank % torch.cuda.device_count())) | |
| torch.cuda.set_device(local_rank) | |
| return rank, world, local_rank | |
| def make_optimizer(decay, no_decay, lr: float, weight_decay: float, cpu: bool): | |
| groups = [ | |
| {"params": decay, "weight_decay": weight_decay}, | |
| {"params": no_decay, "weight_decay": 0.0}, | |
| ] | |
| if not cpu: | |
| try: | |
| return torch.optim.AdamW( | |
| groups, lr=lr, betas=(0.9, 0.95), eps=1e-8, fused=True | |
| ) | |
| except (RuntimeError, TypeError, ValueError) as e: | |
| logger.warning( | |
| "fused AdamW unavailable on this backend (%s); using foreach", e | |
| ) | |
| return torch.optim.AdamW(groups, lr=lr, betas=(0.9, 0.95), eps=1e-8, foreach=True) | |
| def is_main(rank: int) -> bool: | |
| return rank == 0 | |
| def load_training_state(checkpoint_dir: Path) -> dict: | |
| p = checkpoint_dir / "training_state.json" | |
| if p.exists(): | |
| return json.loads(p.read_text(encoding="utf-8")) | |
| return {"stage": None, "global_step": 0} | |
| def lr_at(step: int, total: int, warmup: int, peak: float, floor_ratio: float) -> float: | |
| if step < warmup: | |
| return peak * (step + 1) / max(1, warmup) | |
| if total <= warmup: | |
| return peak | |
| progress = (step - warmup) / max(1, total - warmup) | |
| progress = min(1.0, max(0.0, progress)) | |
| cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) | |
| return peak * (floor_ratio + (1.0 - floor_ratio) * cosine) | |
| def count_corpus_lines(path: Path) -> int: | |
| n = 0 | |
| with path.open("rb") as fh: | |
| for chunk in iter(lambda: fh.read(64 * 1024 * 1024), b""): | |
| n += chunk.count(b"\n") | |
| return n | |
| def estimate_steps_per_epoch( | |
| corpus_lines: int, | |
| avg_tokens_per_doc: float, | |
| seq_len: int, | |
| world: int, | |
| batch: int, | |
| accum: int, | |
| ) -> int: | |
| total_tokens = corpus_lines * avg_tokens_per_doc | |
| sequences = total_tokens / seq_len | |
| per_step = world * batch * accum | |
| return max(1, int(sequences / per_step)) | |
| def gather_full_state(model): | |
| if isinstance(model, FSDP): | |
| cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) | |
| with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, cfg): | |
| return model.state_dict() | |
| return {k: v.detach().cpu() for k, v in model.module.state_dict().items()} | |
| def clip_grads(model, max_norm: float): | |
| if isinstance(model, FSDP): | |
| return model.clip_grad_norm_(max_norm) | |
| return torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm) | |
| def save_checkpoint( | |
| model, | |
| tokenizer, | |
| checkpoint_dir: Path, | |
| stage: GrowthStage, | |
| global_step: int, | |
| rank: int, | |
| max_shard_size: str, | |
| ) -> None: | |
| state = gather_full_state(model) | |
| if is_main(rank): | |
| checkpoint_dir.mkdir(parents=True, exist_ok=True) | |
| for old in list(checkpoint_dir.glob("model*.safetensors*")) + list( | |
| checkpoint_dir.glob("pytorch_model*.bin*") | |
| ): | |
| try: | |
| old.unlink() | |
| except PermissionError: | |
| stale = old.with_name(old.name + ".stale") | |
| try: | |
| old.replace(stale) | |
| except OSError as e: | |
| logger.warning("could not remove or rename %s: %s", old.name, e) | |
| inner = model.module | |
| inner.save_pretrained( | |
| str(checkpoint_dir), | |
| state_dict={k: v.to(torch.float32) for k, v in state.items()}, | |
| safe_serialization=True, | |
| max_shard_size=max_shard_size, | |
| ) | |
| tokenizer.save_pretrained(str(checkpoint_dir)) | |
| (checkpoint_dir / "training_state.json").write_text( | |
| json.dumps({"stage": stage.name, "global_step": global_step}, indent=2), | |
| encoding="utf-8", | |
| ) | |
| logger.info("Saved checkpoint at step %d -> %s", global_step, checkpoint_dir) | |
| dist.barrier() | |
| def cuda_defaults(stage: GrowthStage, world: int) -> tuple[int, int]: | |
| params = stage.n_layer * 12 * stage.n_embd * stage.n_embd | |
| if params < 1e9: | |
| per_gpu_batch = 16 | |
| elif params < 4e9: | |
| per_gpu_batch = 8 | |
| elif params < 10e9: | |
| per_gpu_batch = 4 | |
| else: | |
| per_gpu_batch = 2 | |
| target_tokens_per_step = 1_048_576 | |
| accum = max( | |
| 1, target_tokens_per_step // (world * per_gpu_batch * stage.n_positions) | |
| ) | |
| return per_gpu_batch, accum | |
| def main() -> None: | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--stage", required=True) | |
| parser.add_argument("--epochs", type=float, default=1.0) | |
| parser.add_argument("--max-steps", type=int, default=None) | |
| parser.add_argument("--batch-size", type=int, default=None) | |
| parser.add_argument("--grad-accum", type=int, default=None) | |
| parser.add_argument("--seq-len", type=int, default=None) | |
| parser.add_argument("--lr", type=float, default=3e-4) | |
| parser.add_argument("--min-lr-ratio", type=float, default=0.1) | |
| parser.add_argument("--warmup-steps", type=int, default=200) | |
| parser.add_argument("--weight-decay", type=float, default=0.1) | |
| parser.add_argument("--grad-clip", type=float, default=1.0) | |
| parser.add_argument("--avg-tokens-per-doc", type=float, default=None) | |
| parser.add_argument("--num-workers", type=int, default=4) | |
| parser.add_argument("--log-every", type=int, default=20) | |
| parser.add_argument("--save-every", type=int, default=1000) | |
| parser.add_argument("--max-shard-size", default="5GB") | |
| parser.add_argument("--no-activation-checkpointing", action="store_true") | |
| parser.add_argument("--sharding", choices=["full", "grad_op"], default="full") | |
| parser.add_argument("--seed", type=int, default=1234) | |
| parser.add_argument("--dry-run-cpu", action="store_true") | |
| parser.add_argument("--data-dir", default=str(DATA_DIR)) | |
| parser.add_argument("--checkpoint-root", default=str(CHECKPOINT_ROOT)) | |
| args = parser.parse_args() | |
| cpu = args.dry_run_cpu | |
| rank, world, local_rank = setup_distributed(cpu) | |
| logging.basicConfig( | |
| level=logging.INFO if is_main(rank) else logging.WARNING, | |
| format=f"%(asctime)s %(levelname)s [rank {rank}] %(message)s", | |
| ) | |
| device = torch.device("cpu") if cpu else torch.device("cuda", local_rank) | |
| stage = get_stage(args.stage) | |
| default_batch, default_accum = cuda_defaults(stage, world) | |
| batch_size = args.batch_size or default_batch | |
| grad_accum = args.grad_accum or default_accum | |
| seq_len = args.seq_len or stage.n_positions | |
| checkpoint_dir = Path(args.checkpoint_root) / stage.name | |
| corpus_path = Path(args.data_dir) / "pretrain_corpus.txt" | |
| if not (checkpoint_dir / "config.json").exists(): | |
| raise FileNotFoundError( | |
| f"No checkpoint at {checkpoint_dir}. Grow into {stage.name} first with grow_and_continue.py." | |
| ) | |
| if not corpus_path.exists(): | |
| raise FileNotFoundError(f"{corpus_path} not found") | |
| tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME) | |
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| if is_main(rank): | |
| logger.info("Loading %s from %s", stage.name, checkpoint_dir) | |
| model = RiemannForCausalLM.from_pretrained( | |
| str(checkpoint_dir), torch_dtype=torch.float32 | |
| ) | |
| model.config.use_cache = False | |
| model.gradient_checkpointing = False | |
| if cpu: | |
| model = DDP(model) | |
| else: | |
| wrap_policy = functools.partial( | |
| transformer_auto_wrap_policy, transformer_layer_cls={TransformerBlock} | |
| ) | |
| mixed = MixedPrecision( | |
| param_dtype=torch.bfloat16, | |
| reduce_dtype=torch.bfloat16, | |
| buffer_dtype=torch.bfloat16, | |
| ) | |
| model = FSDP( | |
| model, | |
| auto_wrap_policy=wrap_policy, | |
| mixed_precision=mixed, | |
| sharding_strategy=( | |
| ShardingStrategy.SHARD_GRAD_OP | |
| if args.sharding == "grad_op" | |
| else ShardingStrategy.FULL_SHARD | |
| ), | |
| device_id=local_rank, | |
| use_orig_params=True, | |
| limit_all_gathers=True, | |
| ) | |
| if not args.no_activation_checkpointing: | |
| apply_activation_checkpointing( | |
| model, | |
| checkpoint_wrapper_fn=functools.partial( | |
| checkpoint_wrapper, checkpoint_impl=CheckpointImpl.NO_REENTRANT | |
| ), | |
| check_fn=lambda m: isinstance(m, TransformerBlock), | |
| ) | |
| model.train() | |
| decay, no_decay = [], [] | |
| for name, p in model.named_parameters(): | |
| if not p.requires_grad: | |
| continue | |
| ( | |
| no_decay | |
| if p.ndim < 2 or "layernorm" in name.lower() or "bias" in name | |
| else decay | |
| ).append(p) | |
| optimizer = make_optimizer(decay, no_decay, args.lr, args.weight_decay, cpu) | |
| state = load_training_state(checkpoint_dir) | |
| global_step = int(state.get("global_step", 0)) | |
| corpus_lines = count_corpus_lines(corpus_path) if is_main(rank) else 0 | |
| tokens_per_doc = args.avg_tokens_per_doc | |
| count_path = corpus_path.with_name("corpus_token_count.json") | |
| if tokens_per_doc is None: | |
| if count_path.exists(): | |
| tokens_per_doc = float( | |
| json.loads(count_path.read_text(encoding="utf-8"))["tokens_per_doc"] | |
| ) | |
| else: | |
| tokens_per_doc = 150.0 | |
| if is_main(rank): | |
| logger.warning( | |
| "%s not found; assuming %.0f tokens per document. Run scripts/count_corpus_tokens.py for an exact epoch size.", | |
| count_path.name, | |
| tokens_per_doc, | |
| ) | |
| lines_t = torch.tensor([corpus_lines], device=device) | |
| dist.broadcast(lines_t, src=0) | |
| corpus_lines = int(lines_t.item()) | |
| steps_per_epoch = estimate_steps_per_epoch( | |
| corpus_lines, tokens_per_doc, seq_len, world, batch_size, grad_accum | |
| ) | |
| total_steps = args.max_steps or int(steps_per_epoch * args.epochs) | |
| tokens_per_step = world * batch_size * grad_accum * seq_len | |
| if is_main(rank): | |
| n_params = sum(p.numel() for p in model.parameters()) | |
| logger.info( | |
| "Training %s: world=%d batch=%d accum=%d seq=%d tokens/step=%d, corpus=%d docs, ~%d steps/epoch, target %d steps (%.2f epochs), %.2fB params", | |
| stage.name, | |
| world, | |
| batch_size, | |
| grad_accum, | |
| seq_len, | |
| tokens_per_step, | |
| corpus_lines, | |
| steps_per_epoch, | |
| total_steps, | |
| args.epochs, | |
| n_params / 1e9, | |
| ) | |
| epoch = 0 | |
| accum_loss = torch.zeros((), device=device) | |
| micro = 0 | |
| tokens_seen = 0 | |
| start = time.time() | |
| done = False | |
| while not done: | |
| dataset = PackedTextDataset( | |
| corpus_path, | |
| tokenizer, | |
| seq_len, | |
| shard_id=rank, | |
| num_shards=world, | |
| epoch=epoch, | |
| seed=args.seed, | |
| ) | |
| loader = DataLoader( | |
| dataset, | |
| batch_size=batch_size, | |
| num_workers=args.num_workers, | |
| pin_memory=not cpu, | |
| drop_last=True, | |
| ) | |
| for batch in loader: | |
| batch = batch.to(device, non_blocking=True) | |
| is_last_micro = (micro + 1) % grad_accum == 0 | |
| ctx = model.no_sync() if not is_last_micro else _nullcontext() | |
| with ctx: | |
| out = model(input_ids=batch, labels=batch) | |
| (out.loss / grad_accum).backward() | |
| accum_loss += out.loss.detach() | |
| micro += 1 | |
| tokens_seen += batch.numel() | |
| if not is_last_micro: | |
| continue | |
| lr = lr_at( | |
| global_step, total_steps, args.warmup_steps, args.lr, args.min_lr_ratio | |
| ) | |
| for g in optimizer.param_groups: | |
| g["lr"] = lr | |
| if args.grad_clip > 0: | |
| clip_grads(model, args.grad_clip) | |
| optimizer.step() | |
| optimizer.zero_grad(set_to_none=True) | |
| global_step += 1 | |
| if global_step % args.log_every == 0: | |
| loss_t = accum_loss / (args.log_every * grad_accum) | |
| dist.all_reduce(loss_t, op=dist.ReduceOp.SUM) | |
| loss_t = loss_t / world | |
| if is_main(rank): | |
| elapsed = time.time() - start | |
| tps = (tokens_seen * world) / max(1e-6, elapsed) | |
| remaining = ( | |
| (total_steps - global_step) * tokens_per_step / max(1e-6, tps) | |
| ) | |
| logger.info( | |
| "step %d/%d | loss %.4f | lr %.2e | %.0f tok/s | %.1fh elapsed | ~%.1fh remaining", | |
| global_step, | |
| total_steps, | |
| loss_t.item(), | |
| lr, | |
| tps, | |
| elapsed / 3600, | |
| remaining / 3600, | |
| ) | |
| accum_loss.zero_() | |
| if global_step % args.save_every == 0: | |
| save_checkpoint( | |
| model, | |
| tokenizer, | |
| checkpoint_dir, | |
| stage, | |
| global_step, | |
| rank, | |
| args.max_shard_size, | |
| ) | |
| if global_step >= total_steps: | |
| done = True | |
| break | |
| epoch += 1 | |
| if not done and is_main(rank): | |
| logger.info("Finished corpus pass %d at step %d", epoch, global_step) | |
| save_checkpoint( | |
| model, tokenizer, checkpoint_dir, stage, global_step, rank, args.max_shard_size | |
| ) | |
| if is_main(rank): | |
| logger.info( | |
| "Done. %s at step %d, %.2fB tokens seen, %.2fh wall clock, checkpoint at %s", | |
| stage.name, | |
| global_step, | |
| tokens_seen * world / 1e9, | |
| (time.time() - start) / 3600, | |
| checkpoint_dir, | |
| ) | |
| dist.barrier() | |
| dist.destroy_process_group() | |
| class _nullcontext: | |
| def __enter__(self): | |
| return self | |
| def __exit__(self, *exc): | |
| return False | |
| if __name__ == "__main__": | |
| main() | |