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.
"""
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,