train.py — Training entry point

scripts/train.py is the boundary between a command-line experiment recipe and the training engine. Its responsibility is to select a model, reconcile YAML and CLI configuration, initialize the process and data streams, restore a checkpoint when requested, record the run context, and call train(). It does not contain the update loop, scheduler logic, numerical kernels or checkpoint serialization implementation; those belong to the imported training modules.

Code excerpts below are dedented from their enclosing functions for readability; the displayed statements and ordering are preserved.

Model choice: Kimi K3 or a plumbing check

Before it parses arguments, this file defines SmokeTestLM, a token-local embedding followed by a linear vocabulary head. Given token IDs x with shape (B,T), it returns logits (B,T,V) and an optional flattened cross-entropy loss against targets. Its config is a small dictionary with model_type: smoke_test_lm. There is no attention, sequence mixing or Kimi K3 architecture in this toy model, so its metrics must never be described as model training results.

The build_model() decision starts with the explicit smoke-test switch:

if smoke_test:
    return SmokeTestLM(vocab_size)
if model_config_path is None:
    raise ValueError("Provide --model-config or use --smoke-test")
config = load_config(model_config_path, KimiK3Config)

For Kimi K3, the model YAML is loaded into KimiK3Config. The function rejects a model vocabulary different from the prepared dataset, or a requested seq_len greater than config.max_seq_len. When the model has KDA layers using the reference chunkwise implementation, it also requires seq_len % kda_chunk_size == 0. That alignment restriction is not applied to fused KDA, which handles the final chunk itself, or to an MLA-only model. The function returns a newly constructed KimiK3 without choosing hardware yet.

The CLI and its two configuration layers

The parser requires --data-dir and precisely one of --model-config PATH and --smoke-test. A --training-config YAML may supply the full TrainConfig recipe. Unlike --model-config, it controls optimization, batching, scheduling, precision and runtime settings, not layer dimensions. Omitted CLI recipe flags leave YAML (or TrainConfig defaults) in charge.

The recipe arguments are generated by type groups, preserving their exact hyphenated CLI spelling:

Additionally, --device defaults to cuda when PyTorch detects CUDA and cpu otherwise; it accepts device strings including cuda:<index>. --run-dir, --checkpoint, and --resume select run storage and checkpoint paths. --deterministic enables PyTorch determinism settings; --profile creates a short profiling window with --profile-wait, --profile-warmup, and --profile-active (all default 1, validated by ProfileConfig). These profiling counters measure optimizer attempts, not just successful updates. The parser does not offer a model/backend flag to override kda_implementation; that is part of the model YAML.

The recipe merge matters because a CLI flag intentionally set to False must still override True in YAML:

config = load_config(args.training_config, TrainConfig) if args.training_config else TrainConfig()
overrides = {field.name: getattr(args, field.name) for field in fields(TrainConfig)
             if getattr(args, field.name, None) is not None}
return replace(config, **overrides)

None means “not specified by CLI”; False, zero and nonzero values are explicit overrides. Values not belonging to TrainConfig—such as output paths and the --deterministic switch—are handled separately. Invalid recipe combinations are validated by the config/trainer implementation, not silently fixed by this merge.

A representative small CPU smoke-test invocation from the repository root is below. path/to/prepared-data must be replaced with a real prepared dataset; it is not an included fixture.

python scripts/train.py --data-dir path/to/prepared-data --smoke-test --device cpu --batch-size 1 --seq-len 8 --max-steps 2

This is a genuine training command, not a read-only example: it creates files and performs updates. For inspection without starting training, use python scripts/train.py --help.

Configure process precision and reproducibility

_configure_runtime() runs before model construction. Only cpu and cuda device types are accepted, explicitly requested CUDA must be available, and validate_precision() rejects device/precision incompatibilities. Under --deterministic, the script sets a default CUBLAS_WORKSPACE_CONFIG, enables deterministic PyTorch algorithms, disables cuDNN benchmarking and enables cuDNN deterministic behavior.

random.seed(config.seed)
np.random.seed(config.seed)
torch.manual_seed(config.seed)

Those three seeds fix Python, NumPy and PyTorch default RNGs for this process. Deterministic kernels may not be available for every downstream operation; enabling the switch can cause PyTorch to raise rather than silently substitute an algorithm. Simply supplying a seed is not a claim that independent hardware, software versions or fused kernels reproduce bitwise-identical training. Resume later replaces checkpointed RNG states after the initial setup.

Instantiate independent data streams

_build_data_loaders() creates train and validation datasets from the same prepared directory, checks matching vocabulary sizes, and builds separate TokenDataLoader instances:

train_data = ShardedTokenDataset(data_dir, split="train")
val_data = ShardedTokenDataset(data_dir, split="val")
if train_data.vocab_size != val_data.vocab_size:
    raise ValueError("Training and validation vocab_size must match")
train_loader = TokenDataLoader(train_data, config.batch_size, config.seq_len)
val_loader = TokenDataLoader(val_data, config.batch_size, config.seq_len)

Each loader owns its own cursor; evaluation cannot accidentally advance the train stream. For every batch, the dataset reads B*T+1 contiguous tokens and creates shifted (B,T) int64 CPU tensors. The next batch advances by B*T inputs, keeping the lookahead token as the next first input; data never wraps automatically at the end of a split. The CLI does not shuffle or retokenize the prepared dataset. See dataset.py for cursor semantics and restored state validation.

Restore the training state before opening the run

The main entry point follows a strict order: parse/merge recipe, create optional profile settings, configure runtime, build loaders and model, move the model to the requested device, build the optimizer, initialize TrainState, and only then restore a checkpoint if requested.

optimizer = build_optimizer(model, config)
state = TrainState()
if args.resume:
    state = load_checkpoint(args.resume, model, optimizer, config,
                            vocab_size=train_data.vocab_size,
                            train_loader=train_loader, val_loader=val_loader)

The loader and optimizer must exist before restoration because a continuation includes more than model weights: optimizer groups/state, training counters, the two loader cursors, RNG state and compatible recipe information. The checkpoint loader validates those relationships before allowing continuation. The printed resume line includes the restored step/token counters, current base/saved learning-rate values, train cursor and an RNG-restored statement; that statement is emitted after the loader returns successfully.

Next, start_run() records the effective post-resume configuration in run.json. By default, a new run receives a unique folder beneath results/training/runs/, resolved relative to the current working directory. If a checkpoint beside run.json is resumed, the existing run directory may be reused and a new session appended. Resuming into an existing run checks that the checkpoint belongs to that run and that model/dataset identity is compatible. Otherwise, the run utility may create a new linked run. --checkpoint changes the checkpoint destination, defaulting to last.pt in the run folder (or the resumed checkpoint path when extending an existing run). It refuses paths that would overwrite run metadata or the metrics log.

Diagnostics, artifacts and delegation

_print_run_summary() reports the run/session IDs, manifest/checkpoint/profile paths, device, parameter count, optimizer groups, Muon/AdamW settings, precision, QK-Clip threshold, learning-rate schedule and token budget. The token budget printed here is an arithmetic plan, not measured throughput:

$\text{tokens per nominal update}=\text{batch size}\times\text{sequence length}\times\text{gradient accumulation steps}.$

The actual training loop determines how skipped updates, evaluation, checkpoint saves and profiling affect progress and measured time. The diagnostics distinguish the smoke test explicitly from Kimi K3 training.

Finally, the resource-managed handoff is:

with MetricsWriter(run) as metrics:
    print(f"Metrics: {metrics.path}")
    train(model, optimizer, train_loader, val_loader, config, device,
          state=state, checkpoint_path=run.checkpoint_path, metrics=metrics,
          profile_config=profile_config)

The run component creates/updates run.json; the metrics writer owns metrics.jsonl; the trainer saves checkpoint and optional profiler artifacts. The CLI itself is not a dry-run operation, and --overfit-batch explicitly repeats a fixed batch rather than consuming the ordinary stream. --profile may allocate and write profiler traces. Exact checkpoint format, scheduler application, loss handling, interval semantics and recovery are delegated to trainer.py, run.py and metrics.py.

Verification and boundaries

The inspected test_train_cli.py launches two tiny CPU subprocesses (Kimi K3 and the smoke model), checks exit status, event order, checkpoint format 7, manifests, counters, optimizer groups, recorded MoE diagnostics, model config and a changed real-model weight. Its configuration tests reject incompatible vocab/sequence sizes and cover fused/MLA-only chunk-alignment exceptions (the fused case skips without FLA). test_recipe.py checks representative CLI/YAML overrides, while test_run.py inspects session continuation and run identity. These are assertions inspected in the fixed revision; they were not executed for this documentation pass.