Blame for train.py

from __future__ import annotations
import json
import math
import os
os.environ["NCCL_NVLS_ENABLE"] = "0"
os.environ.setdefault("GLOO_SOCKET_IFNAME", "lo")
import random
import shutil
import signal
import sys
import time
from concurrent.futures import Future, ThreadPoolExecutor
from contextlib import nullcontext
from dataclasses import asdict, dataclass, field, replace
from datetime import timedelta
from pathlib import Path
from typing import Any, Iterator, Optional
import torch
from torch import Tensor, nn
from torch.nn import functional as F
from torch.utils.checkpoint import checkpoint
CONFIG: dict[str, Any] = {
    "run": {
        "stage": "auto",
        "auto_progress": True,
        "auto_start_stage": "pretrain",
        "auto_stop_stage": "eval",
        "run_name": "openmodel1-7b",
        "output_dir": "",
        "output_root": "/data/openmodel-artifacts",
        "seed": 1337,
        "resume": True,
        "max_session_hours": 0.0,
        "stage_time_limits_hours": {
            "pretrain": 11.0 * 24.0,
            "sft": 24.0,
            "simpo": 12.0,
            "eval": 12.0,
        },
        "base_checkpoint": "",
    },
    "model": {
        "variant": "7.0b",
        "vocab_size": 50257,
        "hidden_size": 2560,
        "num_layers": 32,
        "num_attention_heads": 20,
        "num_kv_heads": 4,
        "intermediate_size": 7424,
        "max_position_embeddings": 8192,
        "rope_theta": 1_000_000.0,
        "memory_width": 1536,
        "memory_kernel_size": 7,
        "memory_dilations": (1, 2, 4, 8),
        "layer_pattern": ("global", "global", "global", "recall"),
        "layer_scale_init": 0.5,
        "qk_norm": True,
        "tie_word_embeddings": True,
        "fused_projections": True,
        "dropout": 0.0,
        "gradient_checkpointing": True,
    },
    "training": {
        "seq_len": 4096,
        "micro_batch_size": 1,
        "gradient_accumulation_steps": 16,
        "target_tokens": 0,
        "global_token_accounting": True,
        "chinchilla_tokens_per_parameter": 20.0,
        "chinchilla_extra_fraction": 0.10,
        "sft_target_tokens": 800_000_000,
        "max_steps": 0,
        "learning_rate": 3.0e-4,
        "min_learning_rate": 3.0e-5,
        "warmup_tokens": 600_000_000,
        "weight_decay": 0.1,
        "beta1": 0.9,
        "beta2": 0.95,
        "epsilon": 1.0e-8,
        "grad_clip": 1.0,
        "compute_dtype": "bfloat16",
        "compile_model": True,
        "compile_mode": "default",
        "compile_dynamic": False,
        "native_rms_norm": True,
        "loss_chunk_tokens": 8192,
        "simpo_beta": 2.0,
        "simpo_gamma": 0.5,
        "simpo_steps": 3000,
        "eval_batches": 20,
    },
    "checkpoint": {
        "every_steps": 500,
        "every_minutes": 30.0,
        "remote_every_minutes": 0.0,
        "enable_remote_upload": False,
        "keep_last_local": 1,
        "async_local_save": True,
        "save_optimizer": False,
        "upload_optimizer": False,
        "require_remote_upload": False,
    },
    "data": {
        "tokenizer_name": "EleutherAI/gpt-neox-20b",
        "cache_dir": "/data/huggingface-cache",
        "shuffle_buffer": 10_000,
        "download_retries": 20,
        "stream_retry_attempts": 20,
        "stream_retry_backoff_seconds": 5.0,
        "stream_retry_backoff_max_seconds": 60.0,
        "include_code": False,
        "anneal_data_enabled": True,
        "anneal_start_tokens": -1,
        "anneal_ramp_tokens": 12_000_000_000,
        "pretrain_sources": (
            {
                "dataset_id": "HuggingFaceFW/fineweb-edu",
                "config": "sample-100BT",
                "weight": 0.70,
                "anneal_weight": 0.55,
                "license": "ODC-By / Common Crawl-derived",
            },
            {
                "dataset_id": "HuggingFaceFW/fineweb",
                "config": "sample-100BT",
                "weight": 0.30,
                "anneal_weight": 0.20,
                "license": "ODC-By / Common Crawl-derived",
            },
        ),
        "anneal_sources": (
            {
                "dataset_id": "HuggingFaceTB/finemath",
                "config": "finemath-4plus",
                "text_field": "text",
                "weight": 0.0,
                "anneal_weight": 0.15,
                "license": "ODC-By / Common Crawl-derived",
            },
            {
                "dataset_id": "Cyrile/dataset-the-stack-v2-dedup-sub",
                "config": "Python",
                "text_field": "content",
                "weight": 0.0,
                "anneal_weight": 0.06,
                "license": "per-file license review required",
            },
            {
                "dataset_id": "Cyrile/dataset-the-stack-v2-dedup-sub",
                "config": "JavaScript",
                "text_field": "content",
                "weight": 0.0,
                "anneal_weight": 0.02,
                "license": "per-file license review required",
            },
            {
                "dataset_id": "Cyrile/dataset-the-stack-v2-dedup-sub",
                "config": "Java",
                "text_field": "content",
                "weight": 0.0,
                "anneal_weight": 0.01,
                "license": "per-file license review required",
            },
            {
                "dataset_id": "Cyrile/dataset-the-stack-v2-dedup-sub",
                "config": "C++",
                "text_field": "content",
                "weight": 0.0,
                "anneal_weight": 0.0075,
                "license": "per-file license review required",
            },
            {
                "dataset_id": "Cyrile/dataset-the-stack-v2-dedup-sub",
                "config": "C",
                "text_field": "content",
                "weight": 0.0,
                "anneal_weight": 0.0025,
                "license": "per-file license review required",
            },
        ),
        "code_source": {
            "dataset_id": "Cyrile/dataset-the-stack-v2-dedup-sub",
            "config": "Python",
            "text_field": "content",
            "weight": 0.10,
            "license": "per-file license review required",
        },
        "sft_sources": (
            {
                "dataset_id": "open-r1/OpenR1-Math-220k",
                "messages_field": "messages",
                "weight": 0.55,
                "license": "Apache-2.0",
            },
            {
                "dataset_id": "HuggingFaceTB/smoltalk",
                "messages_field": "messages",
                "weight": 0.45,
                "license": "mixed; audit card",
            },
        ),
        "simpo_sources": (
            {
                "dataset_id": "HuggingFaceH4/ultrafeedback_binarized",
                "split": "train_prefs",
                "weight": 1.0,
                "license": "MIT",
            },
        ),
    },
    "modelscope": {
        "repo_id": "Bc_AI_official/OpenModel-1-Comet-Checkpoints",
        "token": "",
        "private": True,
    },
    "runtime": {
        "distributed": True,
        "backend": "nccl",
        "control_backend": "gloo",
        "expected_world_size": 8,
        "process_group_timeout_minutes": 120,
        "ddp_static_graph": True,
        "ddp_bucket_cap_mb": 256,
        "disable_nvls": True,
    },
}
@dataclass(frozen=True)
class Source:
    dataset_id: str
    config: str | None = None
    split: str = "train"
    text_field: str = "text"
    messages_field: str = "messages"
    weight: float = 1.0
    anneal_weight: float | None = None
    revision: str | None = None
    enabled: bool = True
    license: str = "unknown"
    @property
    def key(self) -> str:
        return "|".join(
            (
                self.dataset_id,
                self.config or "",
                self.split,
                self.revision or "",
            )
        )
@dataclass(frozen=True)
class StageResult:
    stage: str
    checkpoint: Path | None
    complete: bool
    step: int
    tokens_seen: int
@dataclass(frozen=True)
class Settings:
    stage: str = CONFIG["run"]["stage"]
    auto_progress: bool = CONFIG["run"]["auto_progress"]
    auto_start_stage: str = CONFIG["run"]["auto_start_stage"]
    auto_stop_stage: str = CONFIG["run"]["auto_stop_stage"]
    run_name: str = CONFIG["run"]["run_name"]
    output_dir: str = CONFIG["run"]["output_dir"]
    output_root: str = CONFIG["run"]["output_root"]
    seed: int = CONFIG["run"]["seed"]
    resume: bool = CONFIG["run"]["resume"]
    max_session_hours: float = CONFIG["run"]["max_session_hours"]
    stage_time_limits_hours: dict[str, float] = field(
        default_factory=lambda: dict(CONFIG["run"]["stage_time_limits_hours"])
    )
    base_checkpoint: str = CONFIG["run"]["base_checkpoint"]
    variant: str = CONFIG["model"]["variant"]
    vocab_size: int = CONFIG["model"]["vocab_size"]
    hidden_size: int = CONFIG["model"]["hidden_size"]
    num_layers: int = CONFIG["model"]["num_layers"]
    num_attention_heads: int = CONFIG["model"]["num_attention_heads"]
    num_kv_heads: int = CONFIG["model"]["num_kv_heads"]
    intermediate_size: int = CONFIG["model"]["intermediate_size"]
    max_position_embeddings: int = CONFIG["model"]["max_position_embeddings"]
    rope_theta: float = CONFIG["model"]["rope_theta"]
    memory_width: int = CONFIG["model"]["memory_width"]
    memory_kernel_size: int = CONFIG["model"]["memory_kernel_size"]
    memory_dilations: tuple[int, ...] = CONFIG["model"]["memory_dilations"]
    layer_pattern: tuple[str, ...] = CONFIG["model"]["layer_pattern"]
    layer_scale_init: float = CONFIG["model"]["layer_scale_init"]
    qk_norm: bool = CONFIG["model"]["qk_norm"]
    tie_word_embeddings: bool = CONFIG["model"]["tie_word_embeddings"]
    fused_projections: bool = CONFIG["model"]["fused_projections"]
    dropout: float = CONFIG["model"]["dropout"]
    gradient_checkpointing: bool = CONFIG["model"]["gradient_checkpointing"]
    seq_len: int = CONFIG["training"]["seq_len"]
    micro_batch_size: int = CONFIG["training"]["micro_batch_size"]
    gradient_accumulation_steps: int = CONFIG["training"]["gradient_accumulation_steps"]
    target_tokens: int = CONFIG["training"]["target_tokens"]
    global_token_accounting: bool = CONFIG["training"]["global_token_accounting"]
    chinchilla_tokens_per_parameter: float = CONFIG["training"]["chinchilla_tokens_per_parameter"]
    chinchilla_extra_fraction: float = CONFIG["training"]["chinchilla_extra_fraction"]
    sft_target_tokens: int = CONFIG["training"]["sft_target_tokens"]
    max_steps: int = CONFIG["training"]["max_steps"]
    learning_rate: float = CONFIG["training"]["learning_rate"]
    min_learning_rate: float = CONFIG["training"]["min_learning_rate"]
    warmup_tokens: int = CONFIG["training"]["warmup_tokens"]
    weight_decay: float = CONFIG["training"]["weight_decay"]
    beta1: float = CONFIG["training"]["beta1"]
    beta2: float = CONFIG["training"]["beta2"]
    epsilon: float = CONFIG["training"]["epsilon"]
    grad_clip: float = CONFIG["training"]["grad_clip"]
    compute_dtype: str = CONFIG["training"]["compute_dtype"]
    compile_model: bool = CONFIG["training"]["compile_model"]
    compile_mode: str = CONFIG["training"]["compile_mode"]
    compile_dynamic: bool = CONFIG["training"]["compile_dynamic"]
    native_rms_norm: bool = CONFIG["training"]["native_rms_norm"]
    loss_chunk_tokens: int = CONFIG["training"]["loss_chunk_tokens"]
    simpo_beta: float = CONFIG["training"]["simpo_beta"]
    simpo_gamma: float = CONFIG["training"]["simpo_gamma"]
    simpo_steps: int = CONFIG["training"]["simpo_steps"]
    eval_batches: int = CONFIG["training"]["eval_batches"]
    checkpoint_every_steps: int = CONFIG["checkpoint"]["every_steps"]
    checkpoint_every_minutes: float = CONFIG["checkpoint"]["every_minutes"]
    remote_every_minutes: float = CONFIG["checkpoint"]["remote_every_minutes"]
    enable_remote_upload: bool = CONFIG["checkpoint"]["enable_remote_upload"]
    keep_last_local: int = CONFIG["checkpoint"]["keep_last_local"]
    async_local_save: bool = CONFIG["checkpoint"]["async_local_save"]
    save_optimizer: bool = CONFIG["checkpoint"]["save_optimizer"]
    upload_optimizer: bool = CONFIG["checkpoint"]["upload_optimizer"]
    require_remote_upload: bool = CONFIG["checkpoint"]["require_remote_upload"]
    tokenizer_name: str = CONFIG["data"]["tokenizer_name"]
    data_cache_dir: str = CONFIG["data"]["cache_dir"]
    shuffle_buffer: int = CONFIG["data"]["shuffle_buffer"]
    data_download_retries: int = CONFIG["data"]["download_retries"]
    stream_retry_attempts: int = CONFIG["data"]["stream_retry_attempts"]
    stream_retry_backoff_seconds: float = CONFIG["data"]["stream_retry_backoff_seconds"]
    stream_retry_backoff_max_seconds: float = CONFIG["data"]["stream_retry_backoff_max_seconds"]
    include_code: bool = CONFIG["data"]["include_code"]
    anneal_data_enabled: bool = CONFIG["data"]["anneal_data_enabled"]
    anneal_start_tokens: int = CONFIG["data"]["anneal_start_tokens"]
    anneal_ramp_tokens: int = CONFIG["data"]["anneal_ramp_tokens"]
    modelscope_repo_id: str = CONFIG["modelscope"]["repo_id"]
    modelscope_token: str = CONFIG["modelscope"]["token"]
    modelscope_private: bool = CONFIG["modelscope"]["private"]
    distributed: bool = CONFIG["runtime"]["distributed"]
    distributed_backend: str = CONFIG["runtime"]["backend"]
    control_backend: str = CONFIG["runtime"]["control_backend"]
    expected_world_size: int = CONFIG["runtime"]["expected_world_size"]
    process_group_timeout_minutes: int = CONFIG["runtime"]["process_group_timeout_minutes"]
    ddp_static_graph: bool = CONFIG["runtime"]["ddp_static_graph"]
    ddp_bucket_cap_mb: int = CONFIG["runtime"]["ddp_bucket_cap_mb"]
    disable_nvls: bool = CONFIG["runtime"]["disable_nvls"]
    def __post_init__(self) -> None:
        variant = self.variant.lower()
        if variant in {"2.4b", "2.5b"}:
            object.__setattr__(self, "hidden_size", 2560)
            object.__setattr__(self, "num_layers", 32)
            object.__setattr__(self, "num_attention_heads", 20)
            object.__setattr__(self, "num_kv_heads", 4)
            object.__setattr__(self, "intermediate_size", 7424)
            object.__setattr__(self, "memory_width", 1536)
        elif variant in {"3.6b", "4b"}:
            object.__setattr__(self, "hidden_size", 3072)
            object.__setattr__(self, "num_layers", 36)
            object.__setattr__(self, "num_attention_heads", 24)
            object.__setattr__(self, "num_kv_heads", 4)
            object.__setattr__(self, "intermediate_size", 8192)
            object.__setattr__(self, "memory_width", 1792)
        elif variant in {"7b", "7.0b"}:
            object.__setattr__(self, "hidden_size", 4096)
            object.__setattr__(self, "num_layers", 36)
            object.__setattr__(self, "num_attention_heads", 32)
            object.__setattr__(self, "num_kv_heads", 8)
            object.__setattr__(self, "intermediate_size", 12288)
            object.__setattr__(self, "memory_width", 2048)
        else:
            raise ValueError("model.variant must be 2.4b, 3.6b, or 7.0b")
    def model_dict(self) -> dict[str, Any]:
        return {
            "variant": self.variant,
            "vocab_size": self.vocab_size,
            "hidden_size": self.hidden_size,
            "num_hidden_layers": self.num_layers,
            "num_attention_heads": self.num_attention_heads,
            "num_key_value_heads": self.num_kv_heads,
            "intermediate_size": self.intermediate_size,
            "max_position_embeddings": self.max_position_embeddings,
            "rope_theta": self.rope_theta,
            "memory_width": self.memory_width,
            "memory_kernel_size": self.memory_kernel_size,
            "memory_dilations": list(self.memory_dilations),
            "layer_pattern": list(self.layer_pattern),
            "layer_scale_init": self.layer_scale_init,
            "qk_norm": self.qk_norm,
            "tie_word_embeddings": self.tie_word_embeddings,
            "fused_projections": self.fused_projections,
            "dropout": self.dropout,
            "gradient_checkpointing": self.gradient_checkpointing,
        }
    def checkpoint_dir(self) -> Path:
        return Path(self.output_dir)
    def sources(self) -> tuple[Source, ...]:
        source_key = {
            "pretrain": "pretrain_sources",
            "sft": "sft_sources",
            "simpo": "simpo_sources",
        }.get(self.stage)
        if source_key is None:
            return ()
        items = [Source(**item) for item in CONFIG["data"][source_key]]
        if self.stage == "pretrain" and self.anneal_data_enabled:
            items.extend(Source(**item) for item in CONFIG["data"]["anneal_sources"])
        if self.stage == "pretrain" and self.include_code and not self.anneal_data_enabled:
            items.append(Source(**CONFIG["data"]["code_source"]))
        unique: dict[str, Source] = {}
        for source in items:
            unique.setdefault(source.key, source)
        return tuple(unique.values())
    def validate(self) -> None:
        if self.stage not in {"auto", "pretrain", "sft", "simpo", "eval"}:
            raise ValueError("CONFIG['run']['stage'] must be auto, pretrain, sft, simpo, or eval")
        pipeline_stages = ("pretrain", "sft", "simpo", "eval")
        if self.auto_start_stage not in pipeline_stages:
            raise ValueError("CONFIG['run']['auto_start_stage'] is invalid")
        if self.auto_stop_stage not in pipeline_stages:
            raise ValueError("CONFIG['run']['auto_stop_stage'] is invalid")
        if pipeline_stages.index(self.auto_start_stage) > pipeline_stages.index(self.auto_stop_stage):
            raise ValueError("auto_start_stage must not come after auto_stop_stage")
        if self.hidden_size % self.num_attention_heads != 0:
            raise ValueError("hidden_size must divide evenly by num_attention_heads")
        if self.num_attention_heads % self.num_kv_heads != 0:
            raise ValueError("num_attention_heads must divide evenly by num_kv_heads")
        if self.memory_kernel_size < 3 or self.memory_kernel_size % 2 == 0:
            raise ValueError("memory_kernel_size must be an odd integer >= 3")
        if self.seq_len + 1 > self.max_position_embeddings:
            raise ValueError("seq_len + 1 must not exceed max_position_embeddings")
        if self.micro_batch_size < 1 or self.gradient_accumulation_steps < 1:
            raise ValueError("micro_batch_size and gradient_accumulation_steps must be positive")
        if self.loss_chunk_tokens < 1:
            raise ValueError("loss_chunk_tokens must be positive")
        if self.max_session_hours < 0:
            raise ValueError("max_session_hours must be zero or positive")
        unknown_stage_limits = set(self.stage_time_limits_hours) - set(pipeline_stages)
        if unknown_stage_limits:
            raise ValueError(f"stage_time_limits_hours has unknown stages: {sorted(unknown_stage_limits)}")
        if any(hours < 0 for hours in self.stage_time_limits_hours.values()):
            raise ValueError("stage_time_limits_hours values must be zero or positive")
        if self.data_download_retries < 1:
            raise ValueError("data.download_retries must be positive")
        if self.stream_retry_attempts < 1:
            raise ValueError("data.stream_retry_attempts must be positive")
        if self.stream_retry_backoff_seconds < 0:
            raise ValueError("data.stream_retry_backoff_seconds must be zero or positive")
        if self.stream_retry_backoff_max_seconds < self.stream_retry_backoff_seconds:
            raise ValueError("data.stream_retry_backoff_max_seconds must not be smaller than the base backoff")
        if self.anneal_start_tokens < -1:
            raise ValueError("data.anneal_start_tokens must be -1 or non-negative")
        if self.anneal_data_enabled and self.anneal_ramp_tokens < 1:
            raise ValueError("data.anneal_ramp_tokens must be positive when annealing is enabled")
        for source in self.sources():
            if source.weight < 0:
                raise ValueError(f"source weight must be non-negative: {source.key}")
            if source.anneal_weight is not None and source.anneal_weight < 0:
                raise ValueError(f"source anneal_weight must be non-negative: {source.key}")
        if self.stage == "pretrain" and not self.sources():
            raise ValueError("pretrain requires at least one data source")
        if self.distributed_backend not in {"nccl", "gloo"}:
            raise ValueError("CONFIG['runtime']['backend'] must be nccl or gloo")
        if self.control_backend not in {"nccl", "gloo"}:
            raise ValueError("CONFIG['runtime']['control_backend'] must be nccl or gloo")
        if self.expected_world_size < 1:
            raise ValueError("CONFIG['runtime']['expected_world_size'] must be positive")
        if self.ddp_bucket_cap_mb < 1:
            raise ValueError("CONFIG['runtime']['ddp_bucket_cap_mb'] must be positive")
@dataclass
class DistributedContext:
    rank: int = 0
    local_rank: int = 0
    world_size: int = 1
    device: torch.device = torch.device("cpu")
    initialized: bool = False
    control_group: Any | None = None
    control_device: torch.device = torch.device("cpu")
    @property
    def is_main(self) -> bool:
        return self.rank == 0
    def barrier(self) -> None:
        if self.initialized:
            torch.distributed.barrier(group=self.control_group)
    def synchronize_device(self) -> None:
        if self.initialized and self.device.type == "cuda":
            torch.cuda.synchronize(self.device)
    def broadcast_object(self, value: Any) -> Any:
        if not self.initialized:
            return value
        objects = [value if self.is_main else None]
        torch.distributed.broadcast_object_list(objects, src=0, group=self.control_group)
        return objects[0]
    def all_gather_object(self, value: Any) -> list[Any]:
        if not self.initialized:
            return [value]
        gathered: list[Any] = [None for _ in range(self.world_size)]
        torch.distributed.all_gather_object(gathered, value, group=self.control_group)
        return gathered
    def max_flags(self, *flags: bool) -> tuple[bool, ...]:
        if not self.initialized:
            return tuple(flags)
        values = torch.tensor(
            [int(flag) for flag in flags],
            dtype=torch.int32,
            device=self.control_device,
        )
        torch.distributed.all_reduce(
            values,
            op=torch.distributed.ReduceOp.MAX,
            group=self.control_group,
        )
        return tuple(bool(value) for value in values.tolist())
    def close(self) -> None:
        if self.initialized and torch.distributed.is_initialized():
            if self.control_group is not None:
                try:
                    torch.distributed.destroy_process_group(self.control_group)
                except Exception:
                    pass
            torch.distributed.destroy_process_group()
def print0(context: DistributedContext, *values: Any, **kwargs: Any) -> None:
    if context.is_main:
        print(*values, **kwargs)
def init_distributed(settings: Settings) -> DistributedContext:
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    if settings.distributed and world_size != settings.expected_world_size:
        raise RuntimeError(
            "This run is configured for "
            f"{settings.expected_world_size} GPUs, but WORLD_SIZE={world_size}. "
            "Launch with torchrun --standalone --nproc-per-node="
            f"{settings.expected_world_size} openmodel1_pipeline.py"
        )
    if world_size > 1:
        rank = int(os.environ["RANK"])
        local_rank = int(os.environ["LOCAL_RANK"])
        if settings.disable_nvls:
            os.environ["NCCL_NVLS_ENABLE"] = "0"
        if not torch.cuda.is_available():
            raise RuntimeError("Distributed CUDA training requires CUDA on every worker")
        torch.cuda.set_device(local_rank)
        if not torch.distributed.is_initialized():
            init_kwargs = {
                "backend": settings.distributed_backend,
                "timeout": timedelta(minutes=settings.process_group_timeout_minutes),
            }
            try:
                torch.distributed.init_process_group(
                    device_id=torch.device("cuda", local_rank),
                    **init_kwargs,
                )
            except TypeError:
                torch.distributed.init_process_group(**init_kwargs)
        control_group = torch.distributed.new_group(
            backend=settings.control_backend,
            timeout=timedelta(minutes=settings.process_group_timeout_minutes),
        )
        control_device = (
            torch.device("cpu")
            if settings.control_backend == "gloo"
            else torch.device("cuda", local_rank)
        )
        return DistributedContext(
            rank=rank,
            local_rank=local_rank,
            world_size=world_size,
            device=torch.device("cuda", local_rank),
            initialized=True,
            control_group=control_group,
            control_device=control_device,
        )
    if settings.distributed:
        raise RuntimeError(
            "Distributed mode is enabled. Start this file with "
            f"torchrun --standalone --nproc-per-node={settings.expected_world_size} "
            "openmodel1_pipeline.py"
        )
    if not torch.cuda.is_available():
        raise RuntimeError("A CUDA GPU is required for training")
    return DistributedContext(device=torch.device("cuda"))
@dataclass
class Output:
    loss: Optional[Tensor] = None
    logits: Optional[Tensor] = None
    hidden_states: Optional[Tensor] = None
class RMSNorm(nn.Module):
    def __init__(self, size: int, eps: float = 1e-6, use_native: bool = True) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(size))
        self.eps = eps
        self.use_native = use_native and hasattr(F, "rms_norm")
    def forward(self, x: Tensor) -> Tensor:
        dtype = x.dtype
        if self.use_native:
            return F.rms_norm(
                x,
                (x.shape[-1],),
                self.weight.to(dtype),
                self.eps,
            )
        variance = x.float().pow(2).mean(dim=-1, keepdim=True)
        return self.weight.to(dtype) * (x * torch.rsqrt(variance + self.eps).to(dtype))
class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim: int, max_positions: int, theta: float) -> None:
        super().__init__()
        if head_dim % 2:
            raise ValueError("head_dim must be even")
        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
        positions = torch.arange(max_positions).float()
        frequencies = torch.outer(positions, inv_freq)
        self.register_buffer("cos", frequencies.cos(), persistent=False)
        self.register_buffer("sin", frequencies.sin(), persistent=False)
    def forward(self, q: Tensor, k: Tensor, position_ids: Tensor) -> tuple[Tensor, Tensor]:
        cos = self.cos[position_ids].unsqueeze(1).to(q.dtype)
        sin = self.sin[position_ids].unsqueeze(1).to(q.dtype)
        return self.rotate(q, cos, sin), self.rotate(k, cos, sin)
    @staticmethod
    def rotate(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
        even = x[..., ::2]
        odd = x[..., 1::2]
        return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
class GQAAttention(nn.Module):
    def __init__(self, settings: Settings) -> None:
        super().__init__()
        self.heads = settings.num_attention_heads
        self.kv_heads = settings.num_kv_heads
        self.head_dim = settings.hidden_size // settings.num_attention_heads
        self.q_dim = self.heads * self.head_dim
        self.kv_dim = self.kv_heads * self.head_dim
        self.fused_projections = settings.fused_projections
        if self.fused_projections:
            self.qkv_proj = nn.Linear(
                settings.hidden_size,
                self.q_dim + 2 * self.kv_dim,
                bias=False,
            )
        else:
            self.q_proj = nn.Linear(settings.hidden_size, self.q_dim, bias=False)
            self.k_proj = nn.Linear(settings.hidden_size, self.kv_dim, bias=False)
            self.v_proj = nn.Linear(settings.hidden_size, self.kv_dim, bias=False)
        self.o_proj = nn.Linear(self.q_dim, settings.hidden_size, bias=False)
        self.q_norm = RMSNorm(self.head_dim, use_native=settings.native_rms_norm) if settings.qk_norm else nn.Identity()
        self.k_norm = RMSNorm(self.head_dim, use_native=settings.native_rms_norm) if settings.qk_norm else nn.Identity()
        self.rope = RotaryEmbedding(self.head_dim, settings.max_position_embeddings, settings.rope_theta)
    def _load_from_state_dict(
        self,
        state_dict: dict[str, Tensor],
        prefix: str,
        local_metadata: dict[str, Any],
        strict: bool,
        missing_keys: list[str],
        unexpected_keys: list[str],
        error_msgs: list[str],
    ) -> None:
        fused_key = prefix + "qkv_proj.weight"
        legacy_keys = tuple(prefix + name for name in ("q_proj.weight", "k_proj.weight", "v_proj.weight"))
        if self.fused_projections:
            if fused_key not in state_dict and all(key in state_dict for key in legacy_keys):
                state_dict[fused_key] = torch.cat([state_dict[key] for key in legacy_keys], dim=0)
            for key in legacy_keys:
                state_dict.pop(key, None)
        elif fused_key in state_dict and not all(key in state_dict for key in legacy_keys):
            q_weight, k_weight, v_weight = state_dict.pop(fused_key).split(
                (self.q_dim, self.kv_dim, self.kv_dim), dim=0
            )
            state_dict[legacy_keys[0]] = q_weight
            state_dict[legacy_keys[1]] = k_weight
            state_dict[legacy_keys[2]] = v_weight
        super()._load_from_state_dict(
            state_dict,
            prefix,
            local_metadata,
            strict,
            missing_keys,
            unexpected_keys,
            error_msgs,
        )
    def forward(self, x: Tensor, position_ids: Tensor) -> Tensor:
        batch, sequence, _ = x.shape
        if self.fused_projections:
            q, k, v = self.qkv_proj(x).split((self.q_dim, self.kv_dim, self.kv_dim), dim=-1)
        else:
            q = self.q_proj(x)
            k = self.k_proj(x)
            v = self.v_proj(x)
        q = q.view(batch, sequence, self.heads, self.head_dim).transpose(1, 2)
        k = k.view(batch, sequence, self.kv_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch, sequence, self.kv_heads, self.head_dim).transpose(1, 2)
        q, k = self.rope(self.q_norm(q), self.k_norm(k), position_ids)
        if hasattr(F, "scaled_dot_product_attention"):
            try:
                y = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=True, enable_gqa=True)
            except TypeError:
                repeats = self.heads // self.kv_heads
                y = F.scaled_dot_product_attention(
                    q,
                    k.repeat_interleave(repeats, dim=1),
                    v.repeat_interleave(repeats, dim=1),
                    dropout_p=0.0,
                    is_causal=True,
                )
        else:
            repeats = self.heads // self.kv_heads
            k = k.repeat_interleave(repeats, dim=1)
            v = v.repeat_interleave(repeats, dim=1)
            scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
            causal = torch.ones(sequence, sequence, device=x.device, dtype=torch.bool).triu(1)
            scores = scores.masked_fill(causal, torch.finfo(scores.dtype).min)
            y = torch.matmul(scores.softmax(dim=-1), v)
        return self.o_proj(y.transpose(1, 2).contiguous().view(batch, sequence, -1))
class CausalRecallMixer(nn.Module):
    def __init__(self, settings: Settings, dilation: int) -> None:
        super().__init__()
        width = settings.memory_width
        self.width = width
        self.kernel_size = settings.memory_kernel_size
        self.dilation = dilation
        self.in_proj = nn.Linear(settings.hidden_size, 2 * width, bias=False)
        self.depthwise = nn.Conv1d(
            width,
            width,
            kernel_size=self.kernel_size,
            dilation=dilation,
            groups=width,
            bias=False,
        )
        self.out_proj = nn.Linear(width, settings.hidden_size, bias=False)
    def forward(self, x: Tensor) -> Tensor:
        gate, value = self.in_proj(x).chunk(2, dim=-1)
        left = (self.kernel_size - 1) * self.dilation
        value = F.pad(value.transpose(1, 2), (left, 0))
        value = self.depthwise(value).transpose(1, 2)
        return self.out_proj(torch.sigmoid(gate) * value)
class SwiGLU(nn.Module):
    def __init__(self, settings: Settings) -> None:
        super().__init__()
        self.intermediate_size = settings.intermediate_size
        self.fused_projections = settings.fused_projections
        if self.fused_projections:
            self.gate_up = nn.Linear(settings.hidden_size, 2 * settings.intermediate_size, bias=False)
        else:
            self.gate = nn.Linear(settings.hidden_size, settings.intermediate_size, bias=False)
            self.up = nn.Linear(settings.hidden_size, settings.intermediate_size, bias=False)
        self.down = nn.Linear(settings.intermediate_size, settings.hidden_size, bias=False)
    def _load_from_state_dict(
        self,
        state_dict: dict[str, Tensor],
        prefix: str,
        local_metadata: dict[str, Any],
        strict: bool,
        missing_keys: list[str],
        unexpected_keys: list[str],
        error_msgs: list[str],
    ) -> None:
        fused_key = prefix + "gate_up.weight"
        gate_key = prefix + "gate.weight"
        up_key = prefix + "up.weight"
        if self.fused_projections:
            if fused_key not in state_dict and gate_key in state_dict and up_key in state_dict:
                state_dict[fused_key] = torch.cat([state_dict[gate_key], state_dict[up_key]], dim=0)
            state_dict.pop(gate_key, None)
            state_dict.pop(up_key, None)
        elif fused_key in state_dict and not (gate_key in state_dict and up_key in state_dict):
            gate_weight, up_weight = state_dict.pop(fused_key).split(
                (self.intermediate_size, self.intermediate_size), dim=0
            )
            state_dict[gate_key] = gate_weight
            state_dict[up_key] = up_weight
        super()._load_from_state_dict(
            state_dict,
            prefix,
            local_metadata,
            strict,
            missing_keys,
            unexpected_keys,
            error_msgs,
        )
    def forward(self, x: Tensor) -> Tensor:
        if self.fused_projections:
            gate, up = self.gate_up(x).chunk(2, dim=-1)
        else:
            gate = self.gate(x)
            up = self.up(x)
        return self.down(F.silu(gate) * up)
class Block(nn.Module):
    def __init__(self, settings: Settings, index: int) -> None:
        super().__init__()
        self.kind = settings.layer_pattern[index % len(settings.layer_pattern)]
        self.norm1 = RMSNorm(settings.hidden_size, use_native=settings.native_rms_norm)
        self.norm2 = RMSNorm(settings.hidden_size, use_native=settings.native_rms_norm)
        if self.kind == "global":
            self.mixer = GQAAttention(settings)
        else:
            dilation = settings.memory_dilations[index % len(settings.memory_dilations)]
            self.mixer = CausalRecallMixer(settings, dilation)
        self.mlp = SwiGLU(settings)
        self.mixer_scale = nn.Parameter(torch.full((settings.hidden_size,), settings.layer_scale_init))
        self.mlp_scale = nn.Parameter(torch.full((settings.hidden_size,), settings.layer_scale_init))
        self.dropout = settings.dropout
    def forward(self, x: Tensor, position_ids: Tensor) -> Tensor:
        h = self.norm1(x)
        h = self.mixer(h, position_ids) if self.kind == "global" else self.mixer(h)
        if self.dropout:
            h = F.dropout(h, p=self.dropout, training=self.training)
        x = x + h * self.mixer_scale
        h = self.mlp(self.norm2(x))
        if self.dropout:
            h = F.dropout(h, p=self.dropout, training=self.training)
        return x + h * self.mlp_scale
class OpenModel1(nn.Module):
    def __init__(self, settings: Settings) -> None:
        super().__init__()
        self.settings = settings
        self.embed = nn.Embedding(settings.vocab_size, settings.hidden_size)
        self.blocks = nn.ModuleList(Block(settings, i) for i in range(settings.num_layers))
        self.norm = RMSNorm(settings.hidden_size)
        self.lm_head = None if settings.tie_word_embeddings else nn.Linear(settings.hidden_size, settings.vocab_size, bias=False)
        self.gradient_checkpointing = settings.gradient_checkpointing
        self._init_weights()
    def _init_weights(self) -> None:
        base_std = 0.02
        residual_std = base_std / math.sqrt(2.0 * self.settings.num_layers)
        residual_modules: set[nn.Module] = set()
        for block in self.blocks:
            residual_modules.add(block.mlp.down)
            if isinstance(block.mixer, GQAAttention):
                residual_modules.add(block.mixer.o_proj)
            elif isinstance(block.mixer, CausalRecallMixer):
                residual_modules.add(block.mixer.out_proj)
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, 0.0, residual_std if module in residual_modules else base_std)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, 0.0, base_std)
            elif isinstance(module, nn.Conv1d):
                nn.init.normal_(module.weight, 0.0, base_std)
    def _hidden(self, input_ids: Tensor, position_ids: Tensor | None = None) -> Tensor:
        batch, sequence = input_ids.shape
        if sequence > self.settings.max_position_embeddings:
            raise ValueError("input sequence exceeds max_position_embeddings")
        if position_ids is None:
            position_ids = torch.arange(sequence, device=input_ids.device).unsqueeze(0).expand(batch, -1)
        x = self.embed(input_ids)
        for block in self.blocks:
            if self.training and self.gradient_checkpointing:
                x = checkpoint(block, x, position_ids, use_reentrant=False)
            else:
                x = block(x, position_ids)
        return self.norm(x)
    def lm_weight(self) -> Tensor:
        return self.embed.weight if self.lm_head is None else self.lm_head.weight
    def logits(self, hidden: Tensor) -> Tensor:
        return F.linear(hidden, self.lm_weight())
    def chunked_loss(self, hidden: Tensor, labels: Tensor, chunk_tokens: int) -> Tensor:
        flat_hidden = hidden.reshape(-1, hidden.shape[-1])
        flat_labels = labels.reshape(-1)
        valid = flat_labels.ne(-100).sum().clamp_min(1)
        total = flat_hidden.new_zeros((), dtype=torch.float32)
        for start in range(0, flat_hidden.shape[0], chunk_tokens):
            stop = min(start + chunk_tokens, flat_hidden.shape[0])
            logits = F.linear(flat_hidden[start:stop], self.lm_weight())
            total = total + F.cross_entropy(logits.float(), flat_labels[start:stop], ignore_index=-100, reduction="sum")
        return total / valid
    def forward(
        self,
        input_ids: Tensor,
        labels: Tensor | None = None,
        return_logits: bool = False,
        loss_chunk_tokens: int = 4096,
        preference_mask: Tensor | None = None,
    ) -> Output | tuple[Tensor, Tensor]:
        hidden = self._hidden(input_ids)
        if preference_mask is not None:
            return self._sequence_logprob_from_hidden(hidden, input_ids, preference_mask)
        loss = None
        if labels is not None:
            loss = self.chunked_loss(hidden[:, :-1], labels[:, 1:], loss_chunk_tokens)
        logits = self.logits(hidden) if return_logits else None
        return Output(loss=loss, logits=logits, hidden_states=hidden)
    def _sequence_logprob_from_hidden(
        self,
        hidden: Tensor,
        input_ids: Tensor,
        target_mask: Tensor,
        chunk_tokens: int = 2048,
    ) -> tuple[Tensor, Tensor]:
        hidden = hidden[:, :-1]
        targets = input_ids[:, 1:]
        mask = target_mask[:, 1:].float()
        flat_hidden = hidden.reshape(-1, hidden.shape[-1])
        flat_targets = targets.reshape(-1)
        flat_mask = mask.reshape(-1)
        token_scores: list[Tensor] = []
        for start in range(0, flat_hidden.shape[0], chunk_tokens):
            stop = min(start + chunk_tokens, flat_hidden.shape[0])
            logits = F.linear(flat_hidden[start:stop], self.lm_weight()).float()
            log_probs = logits.log_softmax(dim=-1)
            score = log_probs.gather(-1, flat_targets[start:stop, None]).squeeze(-1)
            token_scores.append(score * flat_mask[start:stop])
        scores = torch.cat(token_scores).view(input_ids.shape[0], -1).sum(dim=1)
        counts = mask.sum(dim=1).clamp_min(1.0)
        return scores, counts
    def sequence_logprob(self, input_ids: Tensor, target_mask: Tensor, chunk_tokens: int = 2048) -> tuple[Tensor, Tensor]:
        hidden = self._hidden(input_ids)
        return self._sequence_logprob_from_hidden(hidden, input_ids, target_mask, chunk_tokens)
def estimate_parameters(settings: Settings) -> int:
    d = settings.hidden_size
    head_dim = d // settings.num_attention_heads
    kv_dim = settings.num_kv_heads * head_dim
    global_mixer = d * (d + 2 * kv_dim + d)
    recall_mixer = d * (2 * settings.memory_width) + settings.memory_width * settings.memory_kernel_size + settings.memory_width * d
    mlp = 3 * d * settings.intermediate_size
    total = 0
    for i in range(settings.num_layers):
        mixer = global_mixer if settings.layer_pattern[i % len(settings.layer_pattern)] == "global" else recall_mixer
        total += mixer + mlp + 4 * d
    total += settings.vocab_size * d + d
    if not settings.tie_word_embeddings:
        total += settings.vocab_size * d
    return total
def resolved_pretrain_target_tokens(settings: Settings) -> int:
    chinchilla_plus = math.ceil(
        estimate_parameters(settings)
        * settings.chinchilla_tokens_per_parameter
        * (1.0 + settings.chinchilla_extra_fraction)
    )
    configured_target = int(settings.target_tokens)
    return max(configured_target, chinchilla_plus)
def content_text(value: Any) -> str:
    if isinstance(value, str):
        return value
    if isinstance(value, list):
        parts: list[str] = []
        for item in value:
            if isinstance(item, str):
                parts.append(item)
            elif isinstance(item, dict) and item.get("type") == "text":
                parts.append(str(item.get("text", "")))
        return "".join(parts)
    return str(value)
def render_messages(messages: list[dict[str, Any]]) -> str:
    return "".join(
        f"<|{message.get('role', 'user')}|>\n{content_text(message.get('content', ''))}\n"
        for message in messages
    )
def is_transient_stream_error(exc: BaseException) -> bool:
    if isinstance(exc, (OSError, ConnectionError, TimeoutError)):
        return True
    message = str(exc).lower()
    transient_markers = (
        "bad file descriptor",
        "connection aborted",
        "connection reset",
        "connection refused",
        "connection closed",
        "server disconnected",
        "incomplete read",
        "broken pipe",
        "timed out",
        "temporary failure",
        "too many requests",
        "rate limit",
        "http 429",
        "http 500",
        "http 502",
        "http 503",
        "http 504",
    )
    return any(marker in message for marker in transient_markers)
class Mixture:
    def __init__(
        self,
        settings: Settings,
        sources: tuple[Source, ...],
        rank: int = 0,
        world_size: int = 1,
    ) -> None:
        self.settings = settings
        self.rank = rank
        self.world_size = world_size
        self.sources = [
            source
            for source in sources
            if source.enabled
            and (
                source.weight > 0
                or (
                    settings.anneal_data_enabled
                    and source.anneal_weight is not None
                    and source.anneal_weight > 0
                )
            )
        ]
        if not self.sources:
            raise ValueError("No enabled dataset sources")
        self.rng = random.Random(settings.seed + rank)
        self.counts = {source.key: 0 for source in self.sources}
        self.streams: dict[str, Any] = {}
        self.iterators: dict[str, Iterator[dict[str, Any]]] = {}
        self.tokens_seen = 0
        self.anneal_start_tokens: int | None = (
            settings.anneal_start_tokens if settings.anneal_start_tokens >= 0 else None
        )
        self.rows_iter = self.rows()
    def set_token_position(self, tokens_seen: int, resumed: bool = False) -> None:
        self.tokens_seen = max(0, int(tokens_seen))
        if self.settings.anneal_data_enabled and self.anneal_start_tokens is None:
            self.anneal_start_tokens = self.tokens_seen if resumed else 0
    def anneal_fraction(self) -> float:
        if not self.settings.anneal_data_enabled or self.anneal_start_tokens is None:
            return 0.0
        elapsed = self.tokens_seen - self.anneal_start_tokens
        return max(0.0, min(1.0, elapsed / self.settings.anneal_ramp_tokens))
    def weights(self) -> list[float]:
        fraction = self.anneal_fraction()
        values: list[float] = []
        for source in self.sources:
            target = source.anneal_weight
            if self.settings.anneal_data_enabled and target is not None:
                weight = source.weight + fraction * (target - source.weight)
            else:
                weight = source.weight
            values.append(max(0.0, weight))
        if not values or sum(values) <= 0:
            raise ValueError("The active data mixture has no positive source weight")
        return values
    def report(self) -> dict[str, Any]:
        values = self.weights()
        total = sum(values)
        return {
            "anneal_fraction": self.anneal_fraction(),
            "anneal_start_tokens": self.anneal_start_tokens,
            "anneal_ramp_tokens": self.settings.anneal_ramp_tokens,
            "sources": [
                {
                    "dataset_id": source.dataset_id,
                    "config": source.config,
                    "weight": weight / total,
                }
                for source, weight in zip(self.sources, values)
            ],
        }
    def load(self, source: Source) -> Any:
        from datasets import load_dataset
        kwargs: dict[str, Any] = {
            "split": source.split,
            "streaming": True,
            "cache_dir": self.settings.data_cache_dir,
        }
        if source.config:
            kwargs["name"] = source.config
        if source.revision:
            kwargs["revision"] = source.revision
        try:
            from datasets import DownloadConfig
            kwargs["download_config"] = DownloadConfig(
                max_retries=self.settings.data_download_retries,
            )
        except (ImportError, TypeError):
            pass
        stream = load_dataset(source.dataset_id, **kwargs)
        if self.settings.shuffle_buffer:
            stream = stream.shuffle(
                buffer_size=self.settings.shuffle_buffer,
                seed=self.settings.seed + self.rank,
            )
        return stream
    def iterator(self, source: Source) -> Iterator[dict[str, Any]]:
        if source.key not in self.streams:
            self.streams[source.key] = self.load(source)
        stream = self.streams[source.key]
        offset = self.counts[source.key]
        if offset:
            stream = stream.skip(offset)
        return iter(stream)
    def reset_source(self, source: Source) -> None:
        self.iterators.pop(source.key, None)
        self.streams.pop(source.key, None)
    def next_row(self, source: Source) -> dict[str, Any]:
        for attempt in range(1, self.settings.stream_retry_attempts + 1):
            try:
                iterator = self.iterators.get(source.key)
                if iterator is None:
                    iterator = self.iterator(source)
                    self.iterators[source.key] = iterator
                try:
                    row = next(iterator)
                except StopIteration:
                    self.iterators[source.key] = iter(self.streams[source.key])
                    row = next(self.iterators[source.key])
                self.counts[source.key] += 1
                return row
            except Exception as exc:
                if not is_transient_stream_error(exc) or attempt >= self.settings.stream_retry_attempts:
                    raise
                self.reset_source(source)
                delay = min(
                    self.settings.stream_retry_backoff_max_seconds,
                    self.settings.stream_retry_backoff_seconds * (2 ** (attempt - 1)),
                )
                print(
                    f"[data] rank={self.rank} source={source.key} "
                    f"stream failure {attempt}/{self.settings.stream_retry_attempts}: "
                    f"{exc!r}; reopening in {delay:.1f}s",
                    flush=True,
                )
                if delay:
                    time.sleep(delay)
        raise RuntimeError(f"stream retry loop ended unexpectedly for {source.key}")
    def rows(self) -> Iterator[tuple[Source, dict[str, Any]]]:
        while True:
            source = self.rng.choices(self.sources, weights=self.weights(), k=1)[0]
            yield source, self.next_row(source)
    def state(self) -> dict[str, Any]:
        return {
            "counts": dict(self.counts),
            "rng_state": self.rng.getstate(),
            "tokens_seen": self.tokens_seen,
            "anneal_start_tokens": self.anneal_start_tokens,
        }
    def restore(self, state: dict[str, Any]) -> None:
        counts = state.get("counts", {})
        for source in self.sources:
            value = counts.get(source.key, counts.get(source.dataset_id, 0))
            self.counts[source.key] = int(value)
        if state.get("rng_state") is not None:
            self.rng.setstate(nested_tuple(state["rng_state"]))
        if state.get("tokens_seen") is not None:
            self.tokens_seen = max(0, int(state["tokens_seen"]))
        if state.get("anneal_start_tokens") is not None:
            self.anneal_start_tokens = int(state["anneal_start_tokens"])
        self.iterators.clear()
        self.rows_iter = self.rows()
def nested_tuple(value: Any) -> Any:
    if isinstance(value, list):
        return tuple(nested_tuple(item) for item in value)
    if isinstance(value, tuple):
        return tuple(nested_tuple(item) for item in value)
    return value
class PackedStream:
    def __init__(self, settings: Settings, tokenizer: Any, mixture: Mixture, mode: str) -> None:
        self.settings = settings
        self.tokenizer = tokenizer
        self.mixture = mixture
        self.mode = mode
        self.eos = tokenizer.eos_token_id
        if self.eos is None:
            raise ValueError("Tokenizer must define eos_token_id")
        self.tokens: list[int] = []
        self.labels: list[int] = []
        self.max_buffer_tokens = max(settings.seq_len * 8, 32_768)
        self.rows_iter = mixture.rows_iter
    def __iter__(self) -> "PackedStream":
        return self
    def state(self) -> dict[str, Any]:
        limit = self.max_buffer_tokens
        return {
            "tokens": list(self.tokens[:limit]),
            "labels": list(self.labels[:limit]),
            "truncated": len(self.tokens) > limit,
        }
    def restore(self, state: dict[str, Any]) -> None:
        self.tokens = [int(token) for token in state.get("tokens", [])]
        self.labels = [int(label) for label in state.get("labels", [])]
    def add_pretrain_row(self, source: Source, row: dict[str, Any]) -> None:
        text = row.get(source.text_field)
        if isinstance(text, str) and text.strip():
            self.tokens.extend(self.tokenizer.encode(text, add_special_tokens=False))
            self.tokens.append(self.eos)
            if len(self.tokens) > self.max_buffer_tokens:
                self.tokens = self.tokens[: self.max_buffer_tokens]
    def add_sft_row(self, source: Source, row: dict[str, Any]) -> None:
        messages = row.get(source.messages_field)
        if not isinstance(messages, list) or not messages:
            return
        for message in messages:
            role = str(message.get("role", "user"))
            text = f"<|{role}|>\n{content_text(message.get('content', ''))}\n"
            encoded = self.tokenizer.encode(text, add_special_tokens=False)
            self.tokens.extend(encoded)
            self.labels.extend(encoded if role == "assistant" else [-100] * len(encoded))
            self.tokens.append(self.eos)
            self.labels.append(self.eos if role == "assistant" else -100)
            if len(self.tokens) > self.max_buffer_tokens:
                self.tokens = self.tokens[: self.max_buffer_tokens]
                self.labels = self.labels[: self.max_buffer_tokens]
    def __next__(self) -> dict[str, list[int]]:
        while len(self.tokens) < self.settings.seq_len + 1:
            source, row = next(self.rows_iter)
            if self.mode == "pretrain":
                self.add_pretrain_row(source, row)
            else:
                self.add_sft_row(source, row)
        tokens = self.tokens[: self.settings.seq_len + 1]
        self.tokens = self.tokens[self.settings.seq_len + 1 :]
        if self.mode == "pretrain":
            labels = list(tokens)
        else:
            labels = self.labels[: self.settings.seq_len + 1]
            self.labels = self.labels[self.settings.seq_len + 1 :]
        return {"input_ids": tokens, "labels": labels}
class PreferenceStream:
    def __init__(self, settings: Settings, tokenizer: Any, mixture: Mixture) -> None:
        self.settings = settings
        self.tokenizer = tokenizer
        self.mixture = mixture
        self.eos = tokenizer.eos_token_id
        self.rows_iter = mixture.rows_iter
    def _candidate(self, value: Any) -> tuple[list[int], list[int]]:
        if isinstance(value, list) and value and isinstance(value[0], dict):
            tokens: list[int] = []
            mask: list[int] = []
            for message in value:
                role = str(message.get("role", "user"))
                text = f"<|{role}|>\n{content_text(message.get('content', ''))}\n"
                encoded = self.tokenizer.encode(text, add_special_tokens=False)
                tokens.extend(encoded)
                mask.extend([1 if role == "assistant" else 0] * len(encoded))
                tokens.append(self.eos)
                mask.append(1 if role == "assistant" else 0)
            return tokens[: self.settings.seq_len + 1], mask[: self.settings.seq_len + 1]
        text = content_text(value)
        tokens = self.tokenizer.encode(text, add_special_tokens=False)[: self.settings.seq_len]
        tokens.append(self.eos)
        return tokens, [1] * len(tokens)
    def __iter__(self) -> "PreferenceStream":
        return self
    def state(self) -> dict[str, Any]:
        return {}
    def restore(self, state: dict[str, Any]) -> None:
        del state
    def __next__(self) -> dict[str, Any]:
        while True:
            source, row = next(self.rows_iter)
            chosen = row.get("chosen")
            rejected = row.get("rejected")
            if chosen is None or rejected is None:
                continue
            chosen_ids, chosen_mask = self._candidate(chosen)
            rejected_ids, rejected_mask = self._candidate(rejected)
            if len(chosen_ids) > 1 and len(rejected_ids) > 1:
                return {
                    "chosen_ids": chosen_ids,
                    "chosen_mask": chosen_mask,
                    "rejected_ids": rejected_ids,
                    "rejected_mask": rejected_mask,
                }
def load_tokenizer(settings: Settings) -> Any:
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(
        settings.tokenizer_name,
        use_fast=True,
        cache_dir=settings.data_cache_dir,
    )
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    return tokenizer
def make_stream(
    settings: Settings,
    tokenizer: Any,
    context: DistributedContext | None = None,
) -> tuple[Mixture, Iterator[dict[str, Any]]]:
    rank = context.rank if context is not None else 0
    world_size = context.world_size if context is not None else 1
    mixture = Mixture(settings, settings.sources(), rank=rank, world_size=world_size)
    if settings.stage == "simpo":
        return mixture, PreferenceStream(settings, tokenizer, mixture)
    return mixture, PackedStream(settings, tokenizer, mixture, settings.stage)
def collate_lm(samples: list[dict[str, list[int]]], device: torch.device) -> tuple[Tensor, Tensor]:
    input_ids = torch.tensor([sample["input_ids"] for sample in samples], dtype=torch.long, device=device)
    labels = torch.tensor([sample["labels"] for sample in samples], dtype=torch.long, device=device)
    return input_ids, labels
def pad_preference(samples: list[dict[str, Any]], pad_id: int, device: torch.device) -> tuple[Tensor, Tensor, Tensor, Tensor]:
    def make(key: str) -> tuple[Tensor, Tensor]:
        ids = [sample[f"{key}_ids"] for sample in samples]
        masks = [sample[f"{key}_mask"] for sample in samples]
        length = max(2, max(len(item) for item in ids))
        ids = [item[:length] + [pad_id] * max(0, length - len(item[:length])) for item in ids]
        masks = [item[:length] + [0] * max(0, length - len(item[:length])) for item in masks]
        return (
            torch.tensor(ids, dtype=torch.long, device=device),
            torch.tensor(masks, dtype=torch.bool, device=device),
        )
    chosen_ids, chosen_mask = make("chosen")
    rejected_ids, rejected_mask = make("rejected")
    return chosen_ids, chosen_mask, rejected_ids, rejected_mask
def cpu_copy(value: Any) -> Any:
    if isinstance(value, Tensor):
        return value.detach().cpu()
    if isinstance(value, dict):
        return {key: cpu_copy(item) for key, item in value.items()}
    if isinstance(value, list):
        return [cpu_copy(item) for item in value]
    if isinstance(value, tuple):
        return tuple(cpu_copy(item) for item in value)
    return value
def torch_load(path: Path) -> Any:
    try:
        return torch.load(path, map_location="cpu", weights_only=False)
    except TypeError:
        return torch.load(path, map_location="cpu")
def unwrap_model(model: nn.Module) -> nn.Module:
    while True:
        if hasattr(model, "module"):
            model = model.module
            continue
        if hasattr(model, "_orig_mod"):
            model = model._orig_mod
            continue
        return model
def legacy_projection_names(name: str) -> list[str]:
    if name.endswith("qkv_proj.weight"):
        prefix = name[: -len("qkv_proj.weight")]
        return [prefix + item for item in ("q_proj.weight", "k_proj.weight", "v_proj.weight")]
    if name.endswith("gate_up.weight"):
        prefix = name[: -len("gate_up.weight")]
        return [prefix + item for item in ("gate.weight", "up.weight")]
    return [name]
def optimizer_parameter_names(model: nn.Module, legacy_layout: bool = False) -> list[list[str]]:
    decay: list[str] = []
    no_decay: list[str] = []
    for name, parameter in unwrap_model(model).named_parameters():
        if not parameter.requires_grad:
            continue
        names = legacy_projection_names(name) if legacy_layout else [name]
        target = decay if parameter.ndim >= 2 and "norm" not in name and "scale" not in name else no_decay
        target.extend(names)
    return [decay, no_decay]
def migrate_legacy_optimizer_state(
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    legacy_state: dict[str, Any],
) -> dict[str, Any] | None:
    current = optimizer.state_dict()
    current_names = optimizer_parameter_names(model, legacy_layout=False)
    old_names = optimizer_parameter_names(model, legacy_layout=True)
    old_groups = legacy_state.get("param_groups", [])
    if len(old_groups) != len(current["param_groups"]):
        return None
    old_id_by_name: dict[str, int] = {}
    for group_names, group in zip(old_names, old_groups):
        old_ids = group.get("params", [])
        if len(group_names) != len(old_ids):
            return None
        old_id_by_name.update(zip(group_names, old_ids))
    old_states = legacy_state.get("state", {})
    new_state: dict[int, Any] = {}
    migrated = 0
    for group_names, group in zip(current_names, current["param_groups"]):
        new_ids = group.get("params", [])
        if len(group_names) != len(new_ids):
            return None
        for name, new_id in zip(group_names, new_ids):
            parts = legacy_projection_names(name)
            part_states = [old_states[old_id_by_name[part]] for part in parts if part in old_id_by_name and old_id_by_name[part] in old_states]
            if not part_states:
                continue
            if len(part_states) == 1:
                new_state[new_id] = part_states[0]
                continue
            merged: dict[str, Any] = {}
            for key in set().union(*(state.keys() for state in part_states)):
                values = [state[key] for state in part_states if key in state]
                first = values[0]
                if isinstance(first, Tensor) and first.ndim > 0:
                    try:
                        merged[key] = torch.cat(values, dim=0)
                    except RuntimeError:
                        merged[key] = first
                else:
                    merged[key] = first
            new_state[new_id] = merged
            migrated += 1
    new_groups = []
    for current_group, old_group in zip(current["param_groups"], old_groups):
        group = dict(old_group)
        group["params"] = list(current_group["params"])
        new_groups.append(group)
    if migrated == 0:
        return None
    return {"state": new_state, "param_groups": new_groups}
def rng_state() -> dict[str, Any]:
    state: dict[str, Any] = {"python": random.getstate(), "torch": torch.get_rng_state()}
    try:
        import numpy as np
        state["numpy"] = np.random.get_state()
    except ImportError:
        pass
    if torch.cuda.is_available():
        state["cuda"] = torch.cuda.get_rng_state()
    return state
def restore_rng(state: dict[str, Any]) -> None:
    if state.get("python") is not None:
        random.setstate(nested_tuple(state["python"]))
    try:
        import numpy as np
        if state.get("numpy") is not None:
            value = state["numpy"]
            if isinstance(value, list):
                value = tuple(value)
            np.random.set_state(value)
    except ImportError:
        pass
    if state.get("torch") is not None:
        torch.set_rng_state(state["torch"])
    if state.get("cuda") is not None and torch.cuda.is_available():
        torch.cuda.set_rng_state(state["cuda"])
class Checkpoints:
    def __init__(self, settings: Settings, context: DistributedContext | None = None) -> None:
        self.settings = settings
        self.context = context or DistributedContext()
        self.root = settings.checkpoint_dir()
        self.root.mkdir(parents=True, exist_ok=True)
        self.steps = self.root / "checkpoints"
        self.steps.mkdir(parents=True, exist_ok=True)
        self.latest = self.root / "latest.json"
        self.remote_upload_disabled = False
        self._writer = ThreadPoolExecutor(max_workers=1, thread_name_prefix="checkpoint-writer")
        self._pending_save: Future[None] | None = None
    @staticmethod
    def name(step: int) -> str:
        return f"step_{step:012d}"
    def local_latest(self) -> Path | None:
        if self.latest.exists():
            try:
                payload = json.loads(self.latest.read_text(encoding="utf-8"))
                candidate = self.root / payload["local_checkpoint"]
                if (candidate / "manifest.json").exists():
                    return candidate
            except (OSError, KeyError, json.JSONDecodeError):
                pass
        candidates = sorted(item for item in self.steps.glob("step_*") if (item / "manifest.json").exists())
        return candidates[-1] if candidates else None
    def _prepare_disk_space(self, required_bytes: int, protected: Path | None = None) -> None:
        if not self.context.is_main:
            return
        for stale in self.steps.glob(".*.tmp"):
            if protected is not None and stale == protected:
                continue
            shutil.rmtree(stale, ignore_errors=True)
        margin = 512 * 1024 * 1024
        free = shutil.disk_usage(self.steps).free
        if free >= required_bytes + margin:
            return
        candidates = sorted(
            item for item in self.steps.glob("step_*") if item.is_dir()
        )
        for item in candidates:
            if free >= required_bytes + margin:
                break
            print0(
                self.context,
                f"[checkpoint] removing old local checkpoint {item} to make room for the next save",
            )
            shutil.rmtree(item, ignore_errors=True)
            free = shutil.disk_usage(self.steps).free
    def wait_for_save(self, synchronize: bool = True, raise_errors: bool = True) -> None:
        error = ""
        if self.context.is_main and self._pending_save is not None:
            pending = self._pending_save
            self._pending_save = None
            try:
                pending.result()
            except Exception as exc:
                error = repr(exc)
        if synchronize and self.context.initialized:
            error = self.context.broadcast_object(error if self.context.is_main else None) or ""
        if error:
            if raise_errors:
                raise RuntimeError(f"local checkpoint save failed: {error}")
            print0(self.context, f"[checkpoint] background local save failed: {error}")
    def close(self) -> None:
        try:
            self.wait_for_save(synchronize=False, raise_errors=False)
        finally:
            self._writer.shutdown(wait=True, cancel_futures=False)
    def _write_checkpoint(
        self,
        target: Path,
        temporary: Path,
        model_state: dict[str, Tensor],
        optimizer_state: dict[str, Any] | None,
        metadata: dict[str, Any],
        files: list[str],
    ) -> None:
        from safetensors.torch import save_file
        try:
            save_file(model_state, str(temporary / "model.safetensors"), metadata={"format": "pt"})
            if optimizer_state is not None:
                torch.save(optimizer_state, temporary / "training_state.pt")
            (temporary / "metadata.json").write_text(
                json.dumps(metadata, indent=2, sort_keys=True),
                encoding="utf-8",
            )
            (temporary / "manifest.json").write_text(
                json.dumps(
                    {
                        "step": metadata["step"],
                        "tokens_seen": metadata["tokens_seen"],
                        "files": files,
                    },
                    indent=2,
                ),
                encoding="utf-8",
            )
            if target.exists():
                shutil.rmtree(target)
            temporary.replace(target)
            latest_payload = {
                "run_name": self.settings.run_name,
                "stage": self.settings.stage,
                "step": metadata["step"],
                "tokens_seen": metadata["tokens_seen"],
                "token_count_scope": metadata["token_count_scope"],
                "local_checkpoint": str(target.relative_to(self.root)),
                "files": files,
                "updated_at_unix": time.time(),
            }
            temporary_latest = self.root / ".latest.tmp"
            temporary_latest.write_text(
                json.dumps(latest_payload, indent=2, sort_keys=True),
                encoding="utf-8",
            )
            temporary_latest.replace(self.latest)
            self.trim()
        except Exception:
            shutil.rmtree(temporary, ignore_errors=True)
            raise
    def save(
        self,
        model: nn.Module,
        optimizer: torch.optim.Optimizer | None,
        step: int,
        tokens_seen: int,
        mixture: Mixture,
        stream: Any,
        include_optimizer: bool,
    ) -> Path:
        name = self.name(step)
        target = self.steps / name
        temporary = self.steps / f".{name}.tmp"
        self.wait_for_save(synchronize=True)
        self.context.synchronize_device()
        self.context.barrier()
        local_runtime_state = {
            "mixture": cpu_copy(mixture.state()),
            "stream": cpu_copy(stream.state()) if hasattr(stream, "state") else {},
            "rng": cpu_copy(rng_state()),
        }
        if self.context.is_main:
            if temporary.exists():
                shutil.rmtree(temporary)
            temporary.mkdir(parents=True)
        self.context.barrier()
        runtime_name = f"runtime_state_rank_{self.context.rank}.pt"
        torch.save(local_runtime_state, temporary / runtime_name)
        self.context.barrier()
        runtime_files = [
            f"runtime_state_rank_{rank}.pt" for rank in range(self.context.world_size)
        ]
        files = ["model.safetensors", *runtime_files, "metadata.json", "manifest.json"]
        if include_optimizer and optimizer is not None:
            files.insert(1, "training_state.pt")
        if self.context.is_main:
            model_state = {
                key: value.detach().cpu().contiguous()
                for key, value in unwrap_model(model).state_dict().items()
            }
            optimizer_state = None
            if include_optimizer and optimizer is not None:
                optimizer_state = {
                    "step": step,
                    "tokens_seen": tokens_seen,
                    "token_count_scope": "global" if self.settings.global_token_accounting else "per_rank",
                    "optimizer": cpu_copy(optimizer.state_dict()),
                }
            metadata = {
                "run_name": self.settings.run_name,
                "stage": self.settings.stage,
                "step": step,
                "tokens_seen": tokens_seen,
                "token_count_scope": "global" if self.settings.global_token_accounting else "per_rank",
                "model_parameters": estimate_parameters(self.settings),
                "model_config": self.settings.model_dict(),
                "sources": [asdict(source) for source in self.settings.sources()],
                "data_mix": mixture.report(),
                "created_at_unix": time.time(),
                "files": files,
            }
            required_bytes = sum(
                value.numel() * value.element_size()
                for value in model_state.values()
            )
            self._prepare_disk_space(required_bytes, protected=temporary)
            if self.settings.async_local_save:
                self._pending_save = self._writer.submit(
                    self._write_checkpoint,
                    target,
                    temporary,
                    model_state,
                    optimizer_state,
                    metadata,
                    files,
                )
                print0(self.context, f"[checkpoint] queued local save {target} tokens={tokens_seen}")
            else:
                self._write_checkpoint(
                    target,
                    temporary,
                    model_state,
                    optimizer_state,
                    metadata,
                    files,
                )
                print0(self.context, f"[checkpoint] local save complete {target} tokens={tokens_seen}")
        return target
    def trim(self) -> None:
        keep = max(1, self.settings.keep_last_local)
        candidates = sorted(item for item in self.steps.glob("step_*") if item.is_dir())
        for item in candidates[:-keep]:
            shutil.rmtree(item)
    def hub(self) -> Any | None:
        if not self.context.is_main or not self.settings.modelscope_repo_id:
            return None
        try:
            from modelscope_hub.api import HubApi
        except ImportError:
            from modelscope_hub import HubApi
        token = self.settings.modelscope_token.strip()
        return HubApi(token=token) if token else HubApi()
    def ensure_repo(self, api: Any) -> None:
        try:
            api.create_repo(
                repo_id=self.settings.modelscope_repo_id,
                repo_type="model",
                private=self.settings.modelscope_private,
            )
        except TypeError:
            try:
                api.create_repo(self.settings.modelscope_repo_id, repo_type="model", private=self.settings.modelscope_private)
            except Exception as exc:
                if "exist" not in str(exc).lower():
                    raise
        except Exception as exc:
            if "exist" not in str(exc).lower():
                raise
    def upload_folder(self, api: Any, folder: Path, remote: str, message: str) -> None:
        try:
            api.upload_folder(
                repo_id=self.settings.modelscope_repo_id,
                repo_type="model",
                folder_path=str(folder),
                path_in_repo=remote,
                commit_message=message,
            )
        except TypeError:
            api.upload_folder(
                self.settings.modelscope_repo_id,
                folder_path=str(folder),
                repo_type="model",
                path_in_repo=remote,
                commit_message=message,
            )
    def upload_file(self, api: Any, file_path: Path, remote: str, message: str) -> None:
        try:
            api.upload_file(
                repo_id=self.settings.modelscope_repo_id,
                repo_type="model",
                path_or_fileobj=str(file_path),
                path_in_repo=remote,
                commit_message=message,
            )
        except TypeError:
            api.upload_file(
                self.settings.modelscope_repo_id,
                path_or_fileobj=str(file_path),
                repo_type="model",
                path_in_repo=remote,
                commit_message=message,
            )
    def upload(self, checkpoint: Path, step: int) -> bool:
        if not self.context.is_main:
            return False
        api = self.hub()
        if api is None:
            return False
        self.ensure_repo(api)
        remote_root = f"runs/{self.settings.run_name}/{self.settings.stage}"
        remote_checkpoint = f"{remote_root}/checkpoints/{checkpoint.name}"
        self.upload_folder(api, checkpoint, remote_checkpoint, f"checkpoint {self.settings.run_name} {self.settings.stage} {step}")
        payload = json.loads(self.latest.read_text(encoding="utf-8"))
        payload["remote_checkpoint"] = remote_checkpoint
        temporary = self.root / ".latest.remote.json"
        temporary.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
        try:
            self.upload_file(api, temporary, f"{remote_root}/latest.json", f"latest pointer {self.settings.run_name} {step}")
        finally:
            temporary.unlink(missing_ok=True)
        return True
    def download_file(self, api: Any, remote: str, target: Path) -> Path:
        target.mkdir(parents=True, exist_ok=True)
        try:
            result = api.download_file("%s" % self.settings.modelscope_repo_id, "model", remote, local_dir=str(target))
        except TypeError:
            result = api.download_file(
                repo_id=self.settings.modelscope_repo_id,
                repo_type="model",
                file_path=remote,
                local_dir=str(target),
            )
        candidates = [Path(str(result))] if result else []
        candidates.extend([target / Path(remote).name, target / remote])
        for candidate in candidates:
            if candidate.exists():
                return candidate
        raise FileNotFoundError(remote)
    def remote_latest(self) -> Path | None:
        api = self.hub()
        if api is None:
            return None
        cache = self.root / ".remote_resume"
        if cache.exists():
            shutil.rmtree(cache)
        cache.mkdir(parents=True)
        try:
            latest_path = self.download_file(api, f"runs/{self.settings.run_name}/{self.settings.stage}/latest.json", cache)
            latest_payload = json.loads(latest_path.read_text(encoding="utf-8"))
            remote_checkpoint = latest_payload["remote_checkpoint"]
            step = int(latest_payload["step"])
            local = self.steps / self.name(step)
            if (local / "manifest.json").exists():
                return local
            local.mkdir(parents=True, exist_ok=True)
            for file_name in latest_payload.get("files", []):
                self.download_file(api, f"{remote_checkpoint}/{file_name}", local)
            return local if (local / "manifest.json").exists() else None
        except Exception as exc:
            print(f"[resume] no remote checkpoint found: {exc}")
            return None
    def restore(
        self,
        model: nn.Module,
        optimizer: torch.optim.Optimizer | None,
        mixture: Mixture,
        stream: Any,
    ) -> tuple[int, int, Path | None]:
        from safetensors.torch import load_file
        checkpoint = None
        if self.context.is_main and self.settings.resume:
            checkpoint = self.local_latest()
            if checkpoint is None and self.settings.enable_remote_upload:
                checkpoint = self.remote_latest()
        checkpoint_value = self.context.broadcast_object(str(checkpoint) if checkpoint is not None else None)
        checkpoint = Path(checkpoint_value) if checkpoint_value else None
        if checkpoint is None:
            self.context.barrier()
            return 0, 0, None
        model_state = load_file(str(checkpoint / "model.safetensors"), device="cpu")
        legacy_layout = any(
            key.endswith("q_proj.weight") or key.endswith("gate.weight")
            for key in model_state
        )
        unwrap_model(model).load_state_dict(
            model_state,
            strict=True,
        )
        step = 0
        tokens_seen = 0
        metadata_path = checkpoint / "metadata.json"
        if metadata_path.exists():
            metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
            step = int(metadata.get("step", 0))
            tokens_seen = int(metadata.get("tokens_seen", 0))
            token_count_scope = metadata.get("token_count_scope", "per_rank")
            if self.settings.global_token_accounting and token_count_scope != "global":
                tokens_seen *= self.context.world_size
                print0(
                    self.context,
                    "[resume] converted legacy per-rank token counter to global tokens: "
                    f"{tokens_seen:,}",
                )
        state_path = checkpoint / "training_state.pt"
        if state_path.exists() and optimizer is not None:
            state = torch_load(state_path)
            step = int(state.get("step", 0))
            tokens_seen = int(state.get("tokens_seen", 0))
            token_count_scope = state.get("token_count_scope", "per_rank")
            if self.settings.global_token_accounting and token_count_scope != "global":
                tokens_seen *= self.context.world_size
                print0(
                    self.context,
                    "[resume] converted legacy per-rank token counter to global tokens: "
                    f"{tokens_seen:,}",
                )
            if state.get("optimizer") is not None:
                try:
                    optimizer_state = state["optimizer"]
                    if legacy_layout and self.settings.fused_projections:
                        migrated = migrate_legacy_optimizer_state(model, optimizer, optimizer_state)
                        if migrated is None:
                            raise RuntimeError("legacy optimizer layout could not be migrated")
                        optimizer.load_state_dict(migrated)
                        print0(self.context, "[resume] migrated legacy optimizer moments into packed projections")
                    else:
                        optimizer.load_state_dict(optimizer_state)
                except Exception as exc:
                    print0(
                        self.context,
                        f"[resume] optimizer state was not loaded; starting fresh optimizer ({exc!r})",
                    )
        runtime_rank_path = checkpoint / f"runtime_state_rank_{self.context.rank}.pt"
        runtime_path = checkpoint / "runtime_state.pt"
        if runtime_rank_path.exists():
            local_state = torch_load(runtime_rank_path)
            mixture.restore(local_state["mixture"])
            if hasattr(stream, "restore"):
                stream.restore(local_state.get("stream", {}))
            restore_rng(local_state["rng"])
        elif runtime_path.exists():
            runtime_payload = torch_load(runtime_path)
            states = runtime_payload.get("states", [])
            if states:
                local_state = states[min(self.context.rank, len(states) - 1)]
                mixture.restore(local_state["mixture"])
                if hasattr(stream, "restore"):
                    stream.restore(local_state.get("stream", {}))
                restore_rng(local_state["rng"])
        else:
            data_path = checkpoint / "data_state.pt"
            if data_path.exists():
                mixture.restore(torch_load(data_path))
            rng_path = checkpoint / "rng_state.pt"
            if rng_path.exists():
                restore_rng(torch_load(rng_path))
        self.context.barrier()
        print0(self.context, f"[resume] step={step} tokens={tokens_seen} from {checkpoint}")
        return step, tokens_seen, checkpoint
def seed_everything(seed: int) -> None:
    random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    try:
        import numpy as np
        np.random.seed(seed)
    except ImportError:
        pass
def model_device(context: DistributedContext) -> torch.device:
    if not torch.cuda.is_available():
        raise RuntimeError("A CUDA GPU is required for the intended Blackwell run")
    return context.device
def configure_cuda_backends() -> None:
    torch.set_float32_matmul_precision("high")
    if not torch.cuda.is_available():
        return
    torch.backends.cuda.matmul.allow_tf32 = True
    for name in ("enable_flash_sdp", "enable_mem_efficient_sdp", "enable_math_sdp"):
        setter = getattr(torch.backends.cuda, name, None)
        if setter is not None:
            setter(True)
    reduced_precision = getattr(torch.backends.cuda.matmul, "allow_bf16_reduced_precision_reduction", None)
    if reduced_precision is not None:
        torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = True
def maybe_compile_model(
    model: nn.Module,
    settings: Settings,
    context: DistributedContext,
) -> nn.Module:
    if not settings.compile_model or not hasattr(torch, "compile"):
        return model
    try:
        compiled = torch.compile(
            model,
            mode=settings.compile_mode,
            dynamic=settings.compile_dynamic,
        )
        print0(context, f"[compile] enabled mode={settings.compile_mode} dynamic={settings.compile_dynamic}")
        return compiled
    except TypeError:
        try:
            compiled = torch.compile(model, dynamic=settings.compile_dynamic)
            print0(context, "[compile] enabled with the default compiler mode")
            return compiled
        except Exception as exc:
            print0(context, f"[compile] unavailable; continuing eager: {exc!r}")
    except Exception as exc:
        print0(context, f"[compile] unavailable; continuing eager: {exc!r}")
    return model
def wrap_distributed(model: nn.Module, context: DistributedContext, settings: Settings) -> nn.Module:
    if not context.initialized:
        return model
    from torch.nn.parallel import DistributedDataParallel
    kwargs = {
        "device_ids": [context.local_rank],
        "output_device": context.local_rank,
        "find_unused_parameters": False,
        "gradient_as_bucket_view": True,
        "bucket_cap_mb": settings.ddp_bucket_cap_mb,
    }
    if settings.ddp_static_graph:
        kwargs["static_graph"] = True
    try:
        return DistributedDataParallel(model, forward_sync_buffers=False, **kwargs)
    except TypeError:
        try:
            return DistributedDataParallel(model, broadcast_buffers=False, **kwargs)
        except TypeError:
            kwargs.pop("static_graph", None)
            return DistributedDataParallel(model, broadcast_buffers=False, **kwargs)
def build_optimizer(model: nn.Module, settings: Settings) -> torch.optim.Optimizer:
    decay: list[Tensor] = []
    no_decay: list[Tensor] = []
    for name, parameter in model.named_parameters():
        if not parameter.requires_grad:
            continue
        if parameter.ndim >= 2 and "norm" not in name and "scale" not in name:
            decay.append(parameter)
        else:
            no_decay.append(parameter)
    groups = [
        {"params": decay, "weight_decay": settings.weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]
    kwargs = {
        "lr": settings.learning_rate,
        "betas": (settings.beta1, settings.beta2),
        "eps": settings.epsilon,
    }
    try:
        return torch.optim.AdamW(groups, fused=torch.cuda.is_available(), **kwargs)
    except TypeError:
        return torch.optim.AdamW(groups, **kwargs)
def set_learning_rate(optimizer: torch.optim.Optimizer, settings: Settings, tokens_seen: int) -> float:
    if tokens_seen < settings.warmup_tokens:
        lr = settings.learning_rate * max(tokens_seen, 1) / settings.warmup_tokens
    else:
        progress = min(1.0, max(0.0, (tokens_seen - settings.warmup_tokens) / max(1, settings.target_tokens - settings.warmup_tokens)))
        lr = settings.min_learning_rate + 0.5 * (settings.learning_rate - settings.min_learning_rate) * (1.0 + math.cos(math.pi * progress))
    for group in optimizer.param_groups:
        group["lr"] = lr
    return lr
def save_checkpoint(
    manager: Checkpoints,
    model: nn.Module,
    optimizer: torch.optim.Optimizer | None,
    step: int,
    tokens_seen: int,
    mixture: Mixture,
    stream: Any,
    remote: bool,
) -> Path:
    remote = False
    checkpoint = manager.save(
        model,
        optimizer,
        step,
        tokens_seen,
        mixture,
        stream,
        include_optimizer=manager.settings.save_optimizer,
    )
    if remote:
        manager.wait_for_save(synchronize=True)
        upload_status: dict[str, Any] = {
            "ok": True,
            "error": "",
            "disabled": manager.remote_upload_disabled,
        }
        if manager.context.is_main and not manager.remote_upload_disabled:
            print0(manager.context, f"[checkpoint] uploading {checkpoint.name} to ModelScope")
            try:
                manager.upload(checkpoint, step)
            except Exception as exc:
                upload_status = {"ok": False, "error": repr(exc)}
                print0(manager.context, f"[checkpoint] ModelScope upload failed: {exc}")
                manager.remote_upload_disabled = True
                upload_status["disabled"] = True
        upload_status = manager.context.broadcast_object(upload_status)
        manager.context.barrier()
        if upload_status.get("disabled"):
            manager.remote_upload_disabled = True
        if not upload_status["ok"]:
            if manager.settings.require_remote_upload:
                raise RuntimeError(f"ModelScope checkpoint upload failed: {upload_status['error']}")
            print0(
                manager.context,
                "[checkpoint] continuing with the local checkpoint; "
                "remote uploads are disabled for this process",
            )
    manager.context.barrier()
    print0(manager.context, f"[checkpoint] local={checkpoint} tokens={tokens_seen}")
    return checkpoint
def train_lm(settings: Settings, context: DistributedContext) -> StageResult:
    settings.validate()
    seed_everything(settings.seed)
    device = model_device(context)
    configure_cuda_backends()
    dtype = torch.bfloat16 if settings.compute_dtype.lower() in {"bf16", "bfloat16"} else torch.float16
    tokenizer = load_tokenizer(settings)
    settings = replace(settings, vocab_size=len(tokenizer))
    if settings.stage == "pretrain":
        settings = replace(settings, target_tokens=resolved_pretrain_target_tokens(settings))
    print0(context, f"[model] {settings.run_name} parameters≈{estimate_parameters(settings):,} vocab={settings.vocab_size}")
    print0(context, f"[data] stage={settings.stage} sources={[source.key for source in settings.sources()]}")
    model = OpenModel1(settings).to(device)
    model = maybe_compile_model(model, settings, context)
    model = wrap_distributed(model, context, settings)
    optimizer = build_optimizer(model, settings)
    mixture, stream = make_stream(settings, tokenizer, context)
    manager = Checkpoints(settings, context)
    step, tokens_seen, checkpoint = manager.restore(model, optimizer, mixture, stream)
    mixture.set_token_position(tokens_seen, resumed=checkpoint is not None)
    if checkpoint is None and settings.stage != "pretrain" and settings.base_checkpoint:
        load_weights_only(model, settings.base_checkpoint, context)
    set_learning_rate(optimizer, settings, min(tokens_seen, settings.target_tokens))
    print0(context, f"[data] active_mix={json.dumps(mixture.report(), sort_keys=True)}")
    model.train()
    stop_requested = {"value": False}
    stop_signal: dict[str, int | None] = {"value": None}
    def request_stop(signum: int, _frame: Any) -> None:
        print0(context, f"[signal] received {signum}; saving at the next optimizer boundary")
        stop_requested["value"] = True
        stop_signal["value"] = signum
    old_int = signal.signal(signal.SIGINT, request_stop)
    old_term = signal.signal(signal.SIGTERM, request_stop)
    session_start = time.monotonic()
    session_start_tokens = tokens_seen
    last_local = session_start
    last_remote = session_start
    log_path = settings.checkpoint_dir() / "train.jsonl"
    log_path.parent.mkdir(parents=True, exist_ok=True)
    micro_tokens = settings.micro_batch_size * settings.seq_len
    token_increment = micro_tokens * settings.gradient_accumulation_steps
    if settings.global_token_accounting:
        token_increment *= context.world_size
    last_checkpoint: Path | None = checkpoint
    try:
        while True:
            if settings.max_steps and step >= settings.max_steps:
                break
            if tokens_seen >= settings.target_tokens:
                break
            optimizer.zero_grad(set_to_none=True)
            loss_value = 0.0
            for micro_index in range(settings.gradient_accumulation_steps):
                samples = [next(stream) for _ in range(settings.micro_batch_size)]
                input_ids, labels = collate_lm(samples, device)
                sync_context = (
                    model.no_sync()
                    if context.initialized
                    and micro_index < settings.gradient_accumulation_steps - 1
                    else nullcontext()
                )
                with sync_context:
                    with torch.autocast(device_type="cuda", dtype=dtype):
                        output = model(input_ids, labels=labels, loss_chunk_tokens=settings.loss_chunk_tokens)
                        loss = output.loss / settings.gradient_accumulation_steps
                    loss.backward()
                loss_value += float(output.loss.detach().float()) / settings.gradient_accumulation_steps
            grad_norm = float(torch.nn.utils.clip_grad_norm_(model.parameters(), settings.grad_clip))
            optimizer.step()
            step += 1
            tokens_seen += token_increment
            mixture.set_token_position(tokens_seen)
            lr = set_learning_rate(optimizer, settings, tokens_seen)
            now = time.monotonic()
            if step % 10 == 0:
                elapsed_seconds = max(now - session_start, 1e-6)
                session_tokens_seen = max(0, tokens_seen - session_start_tokens)
                global_session_tokens = (
                    session_tokens_seen
                    if settings.global_token_accounting
                    else session_tokens_seen * context.world_size
                )
                record = {
                    "step": step,
                    "tokens_seen": tokens_seen,
                    "session_tokens_seen": session_tokens_seen,
                    "loss": loss_value,
                    "lr": lr,
                    "grad_norm": grad_norm,
                    "elapsed_hours": elapsed_seconds / 3600.0,
                    "tokens_per_second_per_rank": global_session_tokens / context.world_size / elapsed_seconds,
                    "tokens_per_second_global": global_session_tokens / elapsed_seconds,
                    "anneal_fraction": mixture.anneal_fraction(),
                }
                print0(context, json.dumps(record, sort_keys=True))
                if context.is_main:
                    with log_path.open("a", encoding="utf-8") as log:
                        log.write(json.dumps(record, sort_keys=True) + "\n")
            local_due, remote_due = context.max_flags(
                step % settings.checkpoint_every_steps == 0
                or now - last_local >= settings.checkpoint_every_minutes * 60.0,
                settings.enable_remote_upload
                and settings.remote_every_minutes > 0
                and now - last_remote >= settings.remote_every_minutes * 60.0,
            )
            if local_due or remote_due:
                print0(context, f"[data] checkpoint_mix={json.dumps(mixture.report(), sort_keys=True)}")
                last_checkpoint = save_checkpoint(
                    manager,
                    model,
                    optimizer,
                    step,
                    tokens_seen,
                    mixture,
                    stream,
                    remote=remote_due and settings.enable_remote_upload and bool(settings.modelscope_repo_id),
                )
                last_local = now
                if remote_due:
                    last_remote = now
            if stop_requested["value"]:
                break
            if settings.max_session_hours > 0 and now - session_start >= settings.max_session_hours * 3600.0:
                print0(context, "[session] time budget reached; final checkpoint will be saved")
                break
    finally:
        should_save_final = (
            sys.exc_info()[0] is None
            and stop_signal["value"] in {None, signal.SIGINT}
        )
        try:
            if should_save_final:
                last_checkpoint = save_checkpoint(
                    manager,
                    model,
                    optimizer,
                    step,
                    tokens_seen,
                    mixture,
                    stream,
                    remote=False,
                )
                manager.wait_for_save(synchronize=True)
                print0(context, f"[checkpoint] local save complete {last_checkpoint} tokens={tokens_seen}")
            else:
                print0(context, "[checkpoint] skipped final save because the process group is not known to be healthy")
        finally:
            manager.close()
            signal.signal(signal.SIGINT, old_int)
            signal.signal(signal.SIGTERM, old_term)
    complete = bool(
        (settings.max_steps and step >= settings.max_steps)
        or tokens_seen >= settings.target_tokens
    )
    return StageResult(settings.stage, last_checkpoint, complete, step, tokens_seen)
def load_weights_only(
    model: nn.Module,
    checkpoint: str,
    context: DistributedContext | None = None,
) -> None:
    from safetensors.torch import load_file
    path = Path(checkpoint)
    if path.is_file() and path.name.endswith(".safetensors"):
        weight_file = path
    else:
        weight_file = path / "model.safetensors"
    if not weight_file.exists():
        raise FileNotFoundError(f"No model.safetensors found under {path}")
    unwrap_model(model).load_state_dict(load_file(str(weight_file), device="cpu"), strict=True)
    if context is None or context.is_main:
        print(f"[base] loaded {weight_file}")
def train_simpo(settings: Settings, context: DistributedContext) -> StageResult:
    settings.validate()
    seed_everything(settings.seed)
    device = model_device(context)
    configure_cuda_backends()
    dtype = torch.bfloat16 if settings.compute_dtype.lower() in {"bf16", "bfloat16"} else torch.float16
    tokenizer = load_tokenizer(settings)
    settings = replace(settings, vocab_size=len(tokenizer))
    model = OpenModel1(settings).to(device)
    model = maybe_compile_model(model, settings, context)
    model = wrap_distributed(model, context, settings)
    manager = Checkpoints(settings, context)
    mixture, stream = make_stream(settings, tokenizer, context)
    optimizer = build_optimizer(model, settings)
    step, tokens_seen, checkpoint = manager.restore(model, optimizer, mixture, stream)
    if checkpoint is None and settings.base_checkpoint:
        load_weights_only(model, settings.base_checkpoint, context)
    set_learning_rate(optimizer, settings, min(tokens_seen, settings.target_tokens))
    model.train()
    beta = settings.simpo_beta
    gamma = settings.simpo_gamma
    max_steps = settings.max_steps or settings.simpo_steps
    session_start = time.monotonic()
    last_local = session_start
    last_remote = session_start
    last_checkpoint: Path | None = checkpoint
    stop_requested = {"value": False}
    stop_signal: dict[str, int | None] = {"value": None}
    def request_stop(signum: int, _frame: Any) -> None:
        print0(context, f"[signal] received {signum}; saving at the next optimizer boundary")
        stop_requested["value"] = True
        stop_signal["value"] = signum
    old_int = signal.signal(signal.SIGINT, request_stop)
    old_term = signal.signal(signal.SIGTERM, request_stop)
    try:
        while step < max_steps and not stop_requested["value"]:
            if settings.max_session_hours > 0 and time.monotonic() - session_start >= settings.max_session_hours * 3600.0:
                break
            optimizer.zero_grad(set_to_none=True)
            loss_value = 0.0
            for micro_index in range(settings.gradient_accumulation_steps):
                samples = [next(stream) for _ in range(settings.micro_batch_size)]
                chosen_ids, chosen_mask, rejected_ids, rejected_mask = pad_preference(samples, tokenizer.pad_token_id, device)
                sync_context = (
                    model.no_sync()
                    if context.initialized
                    and micro_index < settings.gradient_accumulation_steps - 1
                    else nullcontext()
                )
                with sync_context:
                    with torch.autocast(device_type="cuda", dtype=dtype):
                        chosen_logp, chosen_count = model(chosen_ids, preference_mask=chosen_mask)
                        rejected_logp, rejected_count = model(rejected_ids, preference_mask=rejected_mask)
                        chosen_reward = beta * chosen_logp / chosen_count
                        rejected_reward = beta * rejected_logp / rejected_count
                        loss = -F.logsigmoid(chosen_reward - rejected_reward - gamma).mean()
                        loss = loss / settings.gradient_accumulation_steps
                    loss.backward()
                loss_value += float(loss.detach().float())
            grad_norm = float(torch.nn.utils.clip_grad_norm_(model.parameters(), settings.grad_clip))
            optimizer.step()
            step += 1
            token_increment = settings.micro_batch_size * settings.gradient_accumulation_steps * settings.seq_len
            if settings.global_token_accounting:
                token_increment *= context.world_size
            tokens_seen += token_increment
            set_learning_rate(optimizer, settings, min(tokens_seen, settings.target_tokens))
            now = time.monotonic()
            if step % 10 == 0:
                record = {"step": step, "loss": loss_value, "grad_norm": grad_norm, "elapsed_hours": (now - session_start) / 3600.0}
                print0(context, json.dumps(record, sort_keys=True))
            local_due, remote_due = context.max_flags(
                step % settings.checkpoint_every_steps == 0
                or now - last_local >= settings.checkpoint_every_minutes * 60.0,
                settings.enable_remote_upload
                and settings.remote_every_minutes > 0
                and now - last_remote >= settings.remote_every_minutes * 60.0,
            )
            if local_due or remote_due:
                last_checkpoint = save_checkpoint(
                    manager,
                    model,
                    optimizer,
                    step,
                    tokens_seen,
                    mixture,
                    stream,
                    remote=remote_due and settings.enable_remote_upload and bool(settings.modelscope_repo_id),
                )
                last_local = now
                if remote_due:
                    last_remote = now
    finally:
        should_save_final = (
            sys.exc_info()[0] is None
            and stop_signal["value"] in {None, signal.SIGINT}
        )
        try:
            if should_save_final:
                last_checkpoint = save_checkpoint(
                    manager,
                    model,
                    optimizer,
                    step,
                    tokens_seen,
                    mixture,
                    stream,
                    remote=False,
                )
                manager.wait_for_save(synchronize=True)
                print0(context, f"[checkpoint] local save complete {last_checkpoint} tokens={tokens_seen}")
            else:
                print0(context, "[checkpoint] skipped final save because the process group is not known to be healthy")
        finally:
            manager.close()
            signal.signal(signal.SIGINT, old_int)
            signal.signal(signal.SIGTERM, old_term)
    return StageResult(settings.stage, last_checkpoint, step >= max_steps, step, tokens_seen)
def run_eval(settings: Settings, context: DistributedContext) -> None:
    if context.initialized and not settings.base_checkpoint:
        raise ValueError("Distributed eval requires CONFIG['run']['base_checkpoint']")
    if context.initialized and not context.is_main:
        context.barrier()
        return
    settings = replace(settings, stage="pretrain")
    settings.validate()
    seed_everything(settings.seed)
    device = model_device(context)
    tokenizer = load_tokenizer(settings)
    settings = replace(settings, vocab_size=len(tokenizer))
    model = OpenModel1(settings).to(device)
    if settings.base_checkpoint:
        load_weights_only(model, settings.base_checkpoint)
    else:
        manager = Checkpoints(settings, context)
        mixture, restore_stream = make_stream(settings, tokenizer, context)
        manager.restore(model, None, mixture, restore_stream)
    model.eval()
    mixture, stream = make_stream(settings, tokenizer, context)
    total_loss = 0.0
    batches = settings.eval_batches
    with torch.no_grad():
        for _ in range(batches):
            sample = next(stream)
            input_ids, labels = collate_lm([sample], device)
            output = model(input_ids, labels=labels, loss_chunk_tokens=settings.loss_chunk_tokens)
            total_loss += float(output.loss)
    mean_loss = total_loss / batches
    print0(context, json.dumps({"eval_batches": batches, "mean_loss": mean_loss, "perplexity": math.exp(min(20.0, mean_loss))}, indent=2))
    if context.initialized:
        context.barrier()
def stage_output_dir(settings: Settings, stage: str) -> str:
    if settings.output_dir:
        if settings.stage == "auto" or settings.auto_progress:
            return str(Path(settings.output_dir) / stage)
        return settings.output_dir
    return str(Path(settings.output_root) / settings.run_name / stage)
def effective_stage_time_limit_hours(settings: Settings, stage: str) -> float:
    limits = [
        value
        for value in (
            settings.max_session_hours,
            settings.stage_time_limits_hours.get(stage, 0.0),
        )
        if value > 0
    ]
    return min(limits, default=0.0)
def settings_for_stage(settings: Settings, stage: str, base_checkpoint: str = "") -> Settings:
    if stage == "pretrain":
        target_tokens = resolved_pretrain_target_tokens(settings)
    elif stage == "sft":
        target_tokens = settings.sft_target_tokens
    else:
        target_tokens = settings.target_tokens
    return replace(
        settings,
        stage=stage,
        output_dir=stage_output_dir(settings, stage),
        base_checkpoint=base_checkpoint,
        max_session_hours=effective_stage_time_limit_hours(settings, stage),
        target_tokens=target_tokens,
    )
def run_auto_pipeline(settings: Settings, context: DistributedContext) -> None:
    settings.validate()
    pipeline_stages = ("pretrain", "sft", "simpo", "eval")
    start_index = pipeline_stages.index(settings.auto_start_stage)
    stop_index = pipeline_stages.index(settings.auto_stop_stage)
    current_checkpoint = settings.base_checkpoint
    if start_index > 0 and not current_checkpoint:
        raise ValueError(
            "CONFIG['run']['base_checkpoint'] is required when auto_start_stage is not pretrain"
        )
    for stage in pipeline_stages[start_index : stop_index + 1]:
        stage_settings = settings_for_stage(settings, stage, current_checkpoint)
        if context.is_main:
            print_plan(stage_settings)
            print(f"[auto] starting stage={stage} checkpoint={current_checkpoint or 'none'}")
        if stage in {"pretrain", "sft"}:
            result = train_lm(stage_settings, context)
        elif stage == "simpo":
            result = train_simpo(stage_settings, context)
        else:
            run_eval(stage_settings, context)
            print0(context, "[auto] evaluation complete ✅")
            return
        if not result.complete:
            print0(
                context,
                f"[auto] session ended during {stage}; "
                "rerun the same command to resume this stage 🔁"
            )
            return
        if result.checkpoint is None:
            raise RuntimeError(f"{stage} completed without producing a checkpoint")
        current_checkpoint = str(result.checkpoint)
        print0(context, f"[auto] stage complete={stage} checkpoint={current_checkpoint} ✅")
    print0(context, "[auto] pipeline complete 🎉")
def print_plan(settings: Settings, context: DistributedContext | None = None) -> None:
    if context is not None and not context.is_main:
        return
    print(
        json.dumps(
            {
                "model": settings.run_name,
                "variant": settings.variant,
                "stage": settings.stage,
                "auto_progress": settings.auto_progress,
                "parameters_estimate": estimate_parameters(settings),
                "hidden_size": settings.hidden_size,
                "layers": settings.num_layers,
                "attention": f"{settings.num_attention_heads}Q/{settings.num_kv_heads}KV GQA",
                "pattern": list(settings.layer_pattern),
                "seq_len": settings.seq_len,
                "micro_batch_size": settings.micro_batch_size,
                "gradient_accumulation_steps": settings.gradient_accumulation_steps,
                "gradient_checkpointing": settings.gradient_checkpointing,
                "fused_projections": settings.fused_projections,
                "compile_model": settings.compile_model,
                "compile_mode": settings.compile_mode,
                "ddp_static_graph": settings.ddp_static_graph,
                "ddp_bucket_cap_mb": settings.ddp_bucket_cap_mb,
                "target_tokens": settings.target_tokens,
                "global_token_accounting": settings.global_token_accounting,
                "tokens_per_parameter_target": (
                    settings.target_tokens / max(1, estimate_parameters(settings))
                ),
                "session_hours": settings.max_session_hours,
                "stage_time_limit_hours": effective_stage_time_limit_hours(settings, settings.stage),
                "disable_nvls": settings.disable_nvls,
                 "control_backend": settings.control_backend,
                 "save_optimizer": settings.save_optimizer,
                 "async_local_save": settings.async_local_save,
                 "enable_remote_upload": settings.enable_remote_upload,
                "require_remote_upload": settings.require_remote_upload,
                "modelscope_repo": settings.modelscope_repo_id or "disabled (local only)",
            },
            indent=2,
        )
    )
def main() -> None:
    settings = Settings()
    settings.validate()
    context = init_distributed(settings)
    try:
        if settings.stage == "auto" or settings.auto_progress:
            run_auto_pipeline(settings, context)
            return
        settings = settings_for_stage(settings, settings.stage, settings.base_checkpoint)
        print_plan(settings, context)
        if settings.stage in {"pretrain", "sft"}:
            train_lm(settings, context)
        elif settings.stage == "simpo":
            train_simpo(settings, context)
        elif settings.stage == "eval":
            run_eval(settings, context)
    finally:
        context.close()
if __name__ == "__main__":
    main()
"""
Pipeline recommendation for the Oct 25, 2026 → Nov 9, 2026 window:
1. Use the 7.0B variant as the quality-first run. It is approximately 7.0B
   parameters with 36 blocks, 4,096 hidden size, 32Q/8KV GQA, and 12,288
   SwiGLU width. Keep 2.4B and 3.6B as cheap ablation variants, not the main
   run, unless the 8-GPU distributed throughput is not ready.
2. Pretrain to the automatic Chinchilla-plus-10% floor at 4,096 context. For
   the current 7.001B parameter configuration this resolves to approximately
   154.030B tokens (22 tokens per parameter). If you later set a larger explicit
   `target_tokens`, the code keeps the larger value. Keep a held-out shard from
   every source.
3. SFT for 800M assistant tokens using OpenR1-Math plus SmolTalk. Audit
   SmolTalk's component licenses before redistribution.
4. SimPO for 3,000 optimizer steps on UltraFeedback. It is reference-free, so
   the policy does not need a second full-size reference model.
5. Run held-out perplexity, then fixed downstream evaluations for ARC,
   HellaSwag, MMLU, GSM8K, IFEval, and HumanEval. Record exact dataset
   revisions and prompt templates in the ModelScope model card.
Time-safe continuous-run procedure:
     - The configured schedule reserves roughly two days before Nov 9: 11 days
       for pretraining, 1 day for SFT, 12 hours for SimPO, and 12 hours for the
       current lightweight evaluation block. Each stage stops cleanly at its
       budget and can resume from its newest checkpoint.
      - `max_session_hours=0.0` means the per-stage schedule in
        `stage_time_limits_hours` is used. Set it above zero only to impose an
        additional hard cap across every stage.
      - Local full checkpoints are saved every 500 steps or 30 minutes.
       - Remote ModelScope upload/resume is disabled for this run. After the
         remote job finishes, upload the completed local checkpoint directory
         manually from the training machine.
      - FineWeb streaming retries transient HTTP/socket failures with
        exponential backoff and reopens the affected source. The legacy Dolma
        dataset script is excluded because current `datasets` releases reject
        dataset scripts.
      - NVLS is disabled from CONFIG for this node's Fabric Manager issue; the
        normal launch command does not need a separate shell override.
      - With `stage="auto"`, the same distributed process advances pretrain →
        SFT → SimPO → eval only after each stage's target is complete.
      - Launch the production run with
        `torchrun --standalone --nproc-per-node=8 openmodel1_pipeline.py`.
        The code uses DDP, rank-aware data streams, rank-zero checkpoint writes,
         and asynchronous rank-zero local checkpoint serialization.
Ablation gate before claiming the custom architecture is better:
    A. dense GQA baseline: change layer_pattern to 32 global blocks;
    B. hybrid default: [global, global, global, recall];
    C. hybrid without residual layer scales.
Run each for the same token budget and compare validation loss, tokens/sec,
peak memory, and fixed downstream benchmark scores. Keep the hybrid only if it
wins on validation loss or quality at comparable throughput; the design is
intentionally easy to ablate without changing the data pipeline.
"""