# ruff: noqa: E402 """Restartable three-phase base pretraining for Hanse-LM.""" import argparse import hashlib import json import math import os import shutil from dataclasses import dataclass from pathlib import Path # Must be set before importing torch. os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True") os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0") import numpy as np import torch from datasets import load_dataset from torch.utils.data import Dataset from tqdm import tqdm from transformers import ( AutoModelForCausalLM, AutoTokenizer, LlamaConfig, LlamaForCausalLM, PreTrainedTokenizerFast, Trainer, TrainingArguments, set_seed, ) from transformers.trainer_utils import get_last_checkpoint SCRIPT_DIR = Path(__file__).resolve().parent PROJECT_DIR = SCRIPT_DIR.parent RUN_NAME = "Hanse-100M-Base-v1" RUN_DIR = PROJECT_DIR / RUN_NAME FINAL_DIR = PROJECT_DIR / f"{RUN_NAME}-FINAL" TOKENIZER_FILE = SCRIPT_DIR / "hanse_tokenizer.json" DATA_ROOT = Path(os.environ.get("HANSE_DATA_DIR", r"C:\hanselm-data")) VOCAB_SIZE = 32_000 BASE_SEQUENCE_LENGTH = 2048 SEED = 42 BATCH_TEXTS = 256 FLUSH_EVERY = 1_000_000 EVAL_TARGET_TOKENS = 2_000_000 CACHE_FORMAT_VERSION = 1 EVAL_HASH_MODULUS = 10_000 EVAL_HASH_BUCKETS = 100 MODEL_PARAMETERS = 99_144_320 TOKENS_PER_STEP = 131_072 SPECIAL_TOKEN_IDS = { "pad_token_id": 0, "bos_token_id": 1, "eos_token_id": 2, "unk_token_id": 3, } @dataclass(frozen=True) class Source: dataset: str config: str revision: str @dataclass(frozen=True) class Phase: name: str requested_tokens: int sequence_length: int micro_batch_size: int gradient_accumulation_steps: int mix: dict[str, float] learning_rate: float warmup_ratio: float decay_type: str | None = None min_lr_ratio: float | None = None @property def tokens_per_step(self) -> int: return self.sequence_length * self.micro_batch_size * self.gradient_accumulation_steps @property def steps(self) -> int: return self.requested_tokens // self.tokens_per_step @property def actual_tokens(self) -> int: return self.steps * self.tokens_per_step @property def source_budgets(self) -> dict[str, int]: return allocate_budgets(self.actual_tokens, self.mix) SOURCES = { "fineweb_de": Source( "HuggingFaceFW/fineweb-2", "deu_Latn", "af9c13333eb981300149d5ca60a8e9d659b276b9", ), "fineweb_edu_en": Source( "HuggingFaceFW/fineweb-edu", "sample-100BT", "87f09149ef4734204d70ed1d046ddc9ca3f2b8f9", ), "finewiki_de": Source( "HuggingFaceFW/finewiki", "de", "8bd13e72e6a002407649b3e898535f42ceb1aeb9", ), "finewiki_en": Source( "HuggingFaceFW/finewiki", "en", "8bd13e72e6a002407649b3e898535f42ceb1aeb9", ), } PHASES = ( Phase( "phase-1", 15_000_000_000, 2048, 2, 32, {"fineweb_de": 0.45, "fineweb_edu_en": 0.40, "finewiki_de": 0.10, "finewiki_en": 0.05}, 6e-4, 0.01, ), Phase( "phase-2", 4_000_000_000, 4096, 1, 32, {"fineweb_de": 0.40, "fineweb_edu_en": 0.35, "finewiki_de": 0.20, "finewiki_en": 0.05}, 3e-4, 0.01, "1-sqrt", 0.10, ), Phase( "phase-3", 1_000_000_000, 8192, 1, 16, {"fineweb_de": 0.35, "fineweb_edu_en": 0.40, "finewiki_de": 0.20, "finewiki_en": 0.05}, 3e-5, 0.0, "cosine", 1 / 3, ), ) EVAL_MIX = PHASES[0].mix assert all(phase.tokens_per_step == TOKENS_PER_STEP for phase in PHASES) def allocate_budgets(total_tokens: int, mix: dict[str, float]) -> dict[str, int]: if not math.isclose(sum(mix.values()), 1.0): raise ValueError("Source fractions must add up to 1.0.") budgets = {source: int(total_tokens * fraction) for source, fraction in mix.items()} budgets[next(iter(budgets))] += total_tokens - sum(budgets.values()) return budgets def sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as file: for block in iter(lambda: file.read(1024 * 1024), b""): digest.update(block) return digest.hexdigest() def phase_dir(phase: Phase) -> Path: return RUN_DIR / phase.name def phase_final_dir(phase: Phase) -> Path: return FINAL_DIR if phase.name == PHASES[-1].name else phase_dir(phase) / f"{phase.name}-final" def cache_paths(name: str) -> tuple[Path, Path]: return DATA_ROOT / f"{name}.bin", DATA_ROOT / f"{name}.json" def create_tokenizer() -> PreTrainedTokenizerFast: tokenizer = PreTrainedTokenizerFast( tokenizer_file=str(TOKENIZER_FILE), pad_token="<|pad|>", bos_token="<|bos|>", eos_token="<|eos|>", unk_token="<|unk|>", additional_special_tokens=[ "<|system|>", "<|user|>", "<|assistant|>", "<|tool|>", "<|tool_result|>", "<|end_of_turn|>", ], model_max_length=BASE_SEQUENCE_LENGTH, ) assert len(tokenizer) == VOCAB_SIZE assert all(getattr(tokenizer, name) == token_id for name, token_id in SPECIAL_TOKEN_IDS.items()) return tokenizer def create_model(tokenizer: PreTrainedTokenizerFast) -> LlamaForCausalLM: model = LlamaForCausalLM( LlamaConfig( vocab_size=len(tokenizer), hidden_size=640, intermediate_size=2048, num_hidden_layers=16, num_attention_heads=10, num_key_value_heads=2, max_position_embeddings=BASE_SEQUENCE_LENGTH, rope_theta=10_000, tie_word_embeddings=True, use_cache=False, pad_token_id=tokenizer.pad_token_id, bos_token_id=tokenizer.bos_token_id, eos_token_id=tokenizer.eos_token_id, ) ) assert model.num_parameters() == MODEL_PARAMETERS return model def is_evaluation_document(example: dict, text: str) -> bool: key = text for field in ("id", "document_id", "doc_id", "url"): value = example.get(field) if isinstance(value, str) and value: key = value break digest = hashlib.blake2b(key.encode(), digest_size=8).digest() return int.from_bytes(digest, "big") % EVAL_HASH_MODULUS < EVAL_HASH_BUCKETS def load_stream(source_name: str, shuffle_seed: int): source = SOURCES[source_name] return load_dataset( source.dataset, source.config, split="train", streaming=True, revision=source.revision, ).shuffle(seed=shuffle_seed, buffer_size=10_000) def cache_metadata( *, cache_name: str, kind: str, target_tokens: int, source_budgets: dict[str, int], tokenizer_hash: str, shuffle_seed: int, sequence_length: int, ) -> dict: return { "cache_format_version": CACHE_FORMAT_VERSION, "cache_name": cache_name, "kind": kind, "target_tokens": target_tokens, "written_tokens": target_tokens, "dtype": "uint16", "sequence_length": sequence_length, "tokenizer_file": str(TOKENIZER_FILE), "tokenizer_sha256": tokenizer_hash, "vocab_size": VOCAB_SIZE, **SPECIAL_TOKEN_IDS, "sources": { name: { "dataset": source.dataset, "config": source.config, "revision": source.revision, } for name, source in SOURCES.items() }, "source_token_budgets": source_budgets, "random_seed": SEED, "shuffle_seed": shuffle_seed, "shuffle_buffer_size": 10_000, "evaluation_partition": { "hash": "blake2b-64", "modulus": EVAL_HASH_MODULUS, "evaluation_buckets": EVAL_HASH_BUCKETS, }, } def token_cache_is_valid(path: Path, metadata_path: Path, expected: dict) -> bool: if not path.is_file() or not metadata_path.is_file(): return False try: metadata = json.loads(metadata_path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError): return False expected_bytes = expected["target_tokens"] * np.dtype(np.uint16).itemsize return metadata == expected and path.stat().st_size == expected_bytes def close_memmap(memmap: np.memmap) -> None: memmap.flush() mmap = getattr(memmap, "_mmap", None) if mmap is not None and not mmap.closed: mmap.close() def append_tokenized(buffer: list[int], texts: list[str], tokenizer: PreTrainedTokenizerFast) -> None: for encoding in tokenizer.backend_tokenizer.encode_batch(texts, add_special_tokens=False): buffer.append(tokenizer.bos_token_id) buffer.extend(encoding.ids) buffer.append(tokenizer.eos_token_id) texts.clear() def build_token_cache( *, path: Path, metadata_path: Path, metadata: dict, tokenizer: PreTrainedTokenizerFast, include_evaluation_documents: bool, ) -> None: if token_cache_is_valid(path, metadata_path, metadata): print(f"[=] Reusing {metadata['cache_name']}: {path}") return partial_path = path.with_suffix(f"{path.suffix}.partial") for old in (partial_path, path, metadata_path): old.unlink(missing_ok=True) target_tokens = metadata["target_tokens"] print( f"[*] Building {metadata['cache_name']}: {target_tokens:,} tokens " f"({target_tokens * 2 / 1_000_000_000:.2f} GB)" ) token_memmap = np.memmap(partial_path, dtype=np.uint16, mode="w+", shape=(target_tokens,)) progress = tqdm(total=target_tokens, desc=metadata["cache_name"], unit="tok") written = 0 def flush(buffer: list[int], limit: int) -> int: nonlocal written amount = min(len(buffer), limit) token_memmap[written : written + amount] = np.asarray(buffer[:amount], dtype=np.uint16) del buffer[:amount] written += amount progress.update(amount) return amount try: for source_index, (source_name, budget) in enumerate(metadata["source_token_budgets"].items()): source_written = 0 token_buffer: list[int] = [] text_batch: list[str] = [] stream = load_stream(source_name, metadata["shuffle_seed"] + source_index) for example in stream: text = example.get("text") if not isinstance(text, str) or not (text := text.strip()): continue if is_evaluation_document(example, text) != include_evaluation_documents: continue text_batch.append(text) if len(text_batch) < BATCH_TEXTS: continue append_tokenized(token_buffer, text_batch, tokenizer) remaining = budget - source_written if len(token_buffer) >= min(FLUSH_EVERY, remaining): source_written += flush(token_buffer, remaining) if source_written == budget: break if source_written < budget and text_batch: append_tokenized(token_buffer, text_batch, tokenizer) if source_written < budget: source_written += flush(token_buffer, budget - source_written) if source_written != budget: raise RuntimeError( f"{source_name} ended after {source_written:,} tokens; expected {budget:,}." ) if written != target_tokens: raise RuntimeError(f"Wrote {written:,} tokens; expected {target_tokens:,}.") close_memmap(token_memmap) token_memmap = None partial_path.replace(path) metadata_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8") except BaseException: if token_memmap is not None: close_memmap(token_memmap) raise finally: progress.close() class MemmapDataset(Dataset): def __init__(self, path: Path, total_tokens: int, sequence_length: int): if total_tokens % sequence_length: raise ValueError("Token cache must contain complete sequences.") self.path = path self.total_tokens = total_tokens self.sequence_length = sequence_length self._data: np.memmap | None = None @property def data(self) -> np.memmap: if self._data is None: self._data = np.memmap( self.path, dtype=np.uint16, mode="r", shape=(self.total_tokens,), ) return self._data def __len__(self) -> int: return self.total_tokens // self.sequence_length def __getitem__(self, index: int) -> dict[str, torch.Tensor]: start = index * self.sequence_length input_ids = torch.tensor(self.data[start : start + self.sequence_length], dtype=torch.long) return {"input_ids": input_ids, "labels": input_ids.clone()} def phase_plan(phase: Phase, tokenizer_hash: str) -> dict: return { "name": phase.name, "requested_tokens": phase.requested_tokens, "actual_tokens": phase.actual_tokens, "steps": phase.steps, "source_token_budgets": phase.source_budgets, "learning_rate": phase.learning_rate, "warmup_ratio": phase.warmup_ratio, "decay_type": phase.decay_type, "min_lr_ratio": phase.min_lr_ratio, "tokenizer_sha256": tokenizer_hash, "sequence_length": phase.sequence_length, "micro_batch_size": phase.micro_batch_size, "gradient_accumulation_steps": phase.gradient_accumulation_steps, "tokens_per_step": phase.tokens_per_step, } def ensure_phase_manifest(phase: Phase, tokenizer_hash: str) -> dict: plan = phase_plan(phase, tokenizer_hash) directory = phase_dir(phase) manifest_path = directory / "phase-manifest.json" directory.mkdir(parents=True, exist_ok=True) if not manifest_path.exists(): manifest_path.write_text(json.dumps(plan, indent=2), encoding="utf-8") return plan existing = json.loads(manifest_path.read_text(encoding="utf-8")) legacy_fields = {"sequence_length", "micro_batch_size", "gradient_accumulation_steps"} legacy_plan = {key: value for key, value in plan.items() if key not in legacy_fields} if phase.name == "phase-1" and existing == legacy_plan: if existing["tokens_per_step"] != TOKENS_PER_STEP: raise RuntimeError("phase-1 legacy manifest has an invalid token batch size.") manifest_path.write_text(json.dumps(plan, indent=2), encoding="utf-8") elif existing != plan: raise RuntimeError(f"{phase.name} manifest differs from this run plan; refusing resume.") return plan def completed_phase_summary(phase: Phase, plan: dict) -> dict | None: marker = phase_dir(phase) / "completed.json" if not marker.is_file() or not phase_final_dir(phase).is_dir(): return None try: summary = json.loads(marker.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError): return None return summary if summary.get("phase_plan") == plan else None def scheduler_args(phase: Phase) -> dict: if phase.decay_type is None: return { "lr_scheduler_type": "constant_with_warmup", "warmup_ratio": phase.warmup_ratio, } warmup_steps = math.ceil(phase.steps * phase.warmup_ratio) return { "lr_scheduler_type": "warmup_stable_decay", "warmup_ratio": phase.warmup_ratio, "lr_scheduler_kwargs": { "num_stable_steps": 0, "num_decay_steps": phase.steps - warmup_steps, "decay_type": phase.decay_type, "min_lr_ratio": phase.min_lr_ratio, }, } def training_arguments(phase: Phase) -> TrainingArguments: return TrainingArguments( output_dir=str(phase_dir(phase)), run_name=f"{RUN_NAME}-{phase.name}", max_steps=phase.steps, per_device_train_batch_size=phase.micro_batch_size, gradient_accumulation_steps=phase.gradient_accumulation_steps, per_device_eval_batch_size=phase.micro_batch_size, learning_rate=phase.learning_rate, weight_decay=0.1, adam_beta1=0.9, adam_beta2=0.95, max_grad_norm=1.0, bf16=True, optim="adamw_torch", logging_steps=10, logging_first_step=True, logging_nan_inf_filter=False, include_num_input_tokens_seen="all", eval_strategy="steps", eval_steps=1000, prediction_loss_only=True, save_steps=1000, save_total_limit=3, report_to=["tensorboard"], logging_dir=str(RUN_DIR / "runs" / phase.name), seed=SEED, data_seed=SEED, **scheduler_args(phase), ) def load_phase_model(phase_index: int, tokenizer: PreTrainedTokenizerFast) -> LlamaForCausalLM: phase = PHASES[phase_index] if phase_index == 0: set_seed(SEED) return create_model(tokenizer) previous = PHASES[phase_index - 1] previous_final = phase_final_dir(previous) if not previous_final.is_dir(): raise RuntimeError(f"Missing completed model for {previous.name}: {previous_final}") config = LlamaConfig.from_pretrained(previous_final) config.max_position_embeddings = phase.sequence_length config.use_cache = False model = LlamaForCausalLM.from_pretrained(previous_final, config=config) assert model.num_parameters() == MODEL_PARAMETERS return model def run_phase( phase_index: int, tokenizer: PreTrainedTokenizerFast, tokenizer_hash: str, eval_path: Path, evaluation_tokens: int, ) -> dict: phase = PHASES[phase_index] plan = ensure_phase_manifest(phase, tokenizer_hash) if completed := completed_phase_summary(phase, plan): print(f"[=] {phase.name} already completed; skipping.") return completed tokenizer.model_max_length = phase.sequence_length cache_path, metadata_path = cache_paths(phase.name) metadata = cache_metadata( cache_name=phase.name, kind="training", target_tokens=phase.actual_tokens, source_budgets=phase.source_budgets, tokenizer_hash=tokenizer_hash, shuffle_seed=SEED + (phase_index + 1) * 10_000, sequence_length=phase.sequence_length, ) build_token_cache( path=cache_path, metadata_path=metadata_path, metadata=metadata, tokenizer=tokenizer, include_evaluation_documents=False, ) args = training_arguments(phase) trainer = Trainer( model=load_phase_model(phase_index, tokenizer), args=args, train_dataset=MemmapDataset(cache_path, phase.actual_tokens, phase.sequence_length), eval_dataset=MemmapDataset(eval_path, evaluation_tokens, phase.sequence_length), ) checkpoint = get_last_checkpoint(args.output_dir) if checkpoint: print(f"[*] Resuming {phase.name} from {checkpoint}") else: print(f"[*] Starting {phase.name} from step zero") torch.cuda.reset_peak_memory_stats() train_result = trainer.train(resume_from_checkpoint=checkpoint) trainer.log_metrics("train", train_result.metrics) trainer.save_metrics("train", train_result.metrics) trainer.save_state() eval_metrics = trainer.evaluate() trainer.log_metrics("eval", eval_metrics) trainer.save_metrics("eval", eval_metrics) final_path = phase_final_dir(phase) trainer.save_model(final_path) tokenizer.save_pretrained(final_path) summary = { "phase_plan": plan, "final_model_path": str(final_path), "peak_vram_gib": torch.cuda.max_memory_allocated() / 1024**3, "train_metrics": train_result.metrics, "eval_metrics": eval_metrics, } (phase_dir(phase) / "completed.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") return summary def evaluation_tokens() -> int: return EVAL_TARGET_TOKENS // BASE_SEQUENCE_LENGTH * BASE_SEQUENCE_LENGTH def preflight(tokenizer: PreTrainedTokenizerFast, tokenizer_hash: str) -> None: DATA_ROOT.mkdir(parents=True, exist_ok=True) if any(part.lower() == "onedrive" for part in DATA_ROOT.resolve().parts): raise RuntimeError(f"HANSE_DATA_DIR must not be inside OneDrive: {DATA_ROOT}") if not torch.cuda.is_available(): raise RuntimeError("No ROCm GPU detected through torch.cuda.") if not torch.cuda.is_bf16_supported(): raise RuntimeError("The detected GPU does not support BF16 training.") model = create_model(tokenizer) print(f"[+] GPU: {torch.cuda.get_device_name(0)}") print(f"[+] Tokenizer: {TOKENIZER_FILE} ({tokenizer_hash})") print(f"[+] Model parameters: {model.num_parameters():,}") del model eval_tokens = evaluation_tokens() assert all(eval_tokens % phase.sequence_length == 0 for phase in PHASES) cache_sizes = { "evaluation": eval_tokens * 2, **{phase.name: phase.actual_tokens * 2 for phase in PHASES}, } for name in cache_sizes: path, _ = cache_paths(name) path.with_suffix(f"{path.suffix}.partial").unlink(missing_ok=True) cache_bytes = sum(cache_sizes.values()) current_cache_bytes = sum( min(path.stat().st_size, expected_bytes) for name, expected_bytes in cache_sizes.items() if (path := cache_paths(name)[0]).is_file() ) additional_bytes = cache_bytes - current_cache_bytes free_bytes = shutil.disk_usage(DATA_ROOT).free print(f"[+] Data root: {DATA_ROOT.resolve()}") print(f"[+] Token caches: {cache_bytes / 1_000_000_000:.2f} GB total") print(f"[+] Additional disk space needed: {additional_bytes / 1_000_000_000:.2f} GB") print(f"[+] Available disk space: {free_bytes / 1_000_000_000:.2f} GB") print("[*] Phase plan:") for phase in PHASES: print( f" {phase.name}: {phase.requested_tokens // 1_000_000_000}B, " f"context={phase.sequence_length}, micro={phase.micro_batch_size}, " f"accum={phase.gradient_accumulation_steps}, tokens/step={phase.tokens_per_step}" ) print(f" evaluation: {eval_tokens:,} tokens") if free_bytes < additional_bytes: raise RuntimeError( "Insufficient disk space for token caches. " "Set HANSE_DATA_DIR to a larger non-OneDrive volume." ) def verify_final_model() -> None: tokenizer = AutoTokenizer.from_pretrained(FINAL_DIR, local_files_only=True) model, loading_info = AutoModelForCausalLM.from_pretrained( FINAL_DIR, local_files_only=True, output_loading_info=True, ) assert not any( loading_info[key] for key in ("missing_keys", "unexpected_keys", "mismatched_keys") ) assert len(tokenizer) == VOCAB_SIZE assert model.num_parameters() == MODEL_PARAMETERS assert model.config.max_position_embeddings == PHASES[-1].sequence_length assert tokenizer.model_max_length == PHASES[-1].sequence_length assert model.get_input_embeddings().weight.data_ptr() == model.get_output_embeddings().weight.data_ptr() model.config.use_cache = True inputs = tokenizer("Hanse", return_tensors="pt") inputs.pop("token_type_ids", None) with torch.inference_mode(): output = model.generate(**inputs, max_new_tokens=1, do_sample=False) assert output.shape[1] == inputs["input_ids"].shape[1] + 1 def print_final_summary(summaries: list[dict]) -> None: print("[+] Base pretraining complete:") for summary in summaries: plan = summary["phase_plan"] print( f" {plan['name']}: {plan['actual_tokens']:,} tokens, " f"train_loss={summary['train_metrics'].get('train_loss')}, " f"eval_loss={summary['eval_metrics'].get('eval_loss')}, " f"peak_vram={summary['peak_vram_gib']:.2f} GiB" ) print(f" total training tokens: {sum(s['phase_plan']['actual_tokens'] for s in summaries):,}") print(f" final output: {FINAL_DIR}") def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--dry-run", action="store_true") args = parser.parse_args() if not TOKENIZER_FILE.is_file(): raise FileNotFoundError(f"Tokenizer not found: {TOKENIZER_FILE}") tokenizer = create_tokenizer() tokenizer_hash = sha256_file(TOKENIZER_FILE) preflight(tokenizer, tokenizer_hash) if args.dry_run: print("[+] Dry run complete; no data was downloaded or training started.") return eval_tokens = evaluation_tokens() eval_path, eval_metadata_path = cache_paths("evaluation") build_token_cache( path=eval_path, metadata_path=eval_metadata_path, metadata=cache_metadata( cache_name="evaluation", kind="evaluation", target_tokens=eval_tokens, source_budgets=allocate_budgets(eval_tokens, EVAL_MIX), tokenizer_hash=tokenizer_hash, shuffle_seed=SEED, sequence_length=BASE_SEQUENCE_LENGTH, ), tokenizer=tokenizer, include_evaluation_documents=True, ) summaries = [ run_phase(index, tokenizer, tokenizer_hash, eval_path, eval_tokens) for index in range(len(PHASES)) ] verify_final_model() print_final_summary(summaries) if __name__ == "__main__": main()