| |
| """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 |
|
|
| |
| 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() |
|
|
