| # Convert a DecoderStack-medium (d24) capture into a nanochat checkpoint. | |
| # | |
| # DecoderStack writes its own two-file capture (see § Checkpoint capture in | |
| # decoderstack_medium_pt-sft.py): | |
| # | |
| # model_stepNNNNNN.pt {step, code, weights: {name: tensor}} | |
| # bf16 live weights + fp32 scalars, banked over layers. | |
| # optim_stepNNNNNN.pt {step, t_step, state: {"name.attr": tensor}} | |
| # mantissa / frst_mntm / scnd_mntm / exp_avg / exp_avg_sq, | |
| # all-gathered to full size (world-agnostic). | |
| # | |
| # nanochat wants one flat state_dict per capture, in its own module-path key | |
| # names, next to a meta_NNNNNN.json: | |
| # | |
| # <out_dir>/model_NNNNNN.pt torch.save(model.state_dict()) | |
| # <out_dir>/meta_NNNNNN.json {"step", "val_bpb", "model_config", ...} | |
| # | |
| # This script does that translation. The architectures are the same model -- | |
| # DecoderStack-medium is a flattened port of nanochat d24 -- so every tensor has | |
| # a home and nothing is reshaped or transposed: DecoderStack's banks index the | |
| # layer on dim 0 and each slice already uses F.linear's (out, in) convention. | |
| # The only real work is naming, unbanking, and dtype. | |
| # | |
| # DTYPE, and why the optimizer file matters | |
| # ----------------------------------------- | |
| # nanochat holds fp32 master weights for everything except the two embedding | |
| # tables, which init_weights() casts to COMPUTE_DTYPE (bf16). DecoderStack holds | |
| # bf16 LIVE weights everywhere plus a uint16 `mantissa` in the optimizer file -- | |
| # the fp32 master's bit pattern is (live_bf16_bits << 16) | mantissa. So: | |
| # | |
| # nanochat fp32 params <- fp32 master = live + mantissa (needs --optim) | |
| # nanochat bf16 params <- live bf16 as-is (mantissa dropped, | |
| # which is correct: nanochat's bf16 embedding IS its | |
| # master, it has no lower bits to carry) | |
| # | |
| # Without --optim the fp32 params are filled with the bf16 live values upcast to | |
| # fp32 -- the right dtype, but only bf16 precision. That is what you get if you | |
| # only downloaded the model file, and it is fine for eval: on a held-out English | |
| # paragraph the d24 step-5568 capture scores 2.4828 nats/token with the masters | |
| # and 2.4833 without (bpb 0.7252 vs 0.7254). Pass --optim when you want the | |
| # exact master anyway -- it is a bit-exact reconstruction, not an approximation. | |
| # | |
| # OPTIMIZER STATE | |
| # --------------- | |
| # --world-size N additionally writes optim_NNNNNN_rank{0..N-1}.pt, the ZeRO-2 | |
| # shards nanochat's DistMuonAdamW expects, so chat_sft can warm-start its | |
| # optimizer instead of printing "starting with fresh optimizer (slightly worse)". | |
| # It is optional: SFT is correct without it. Pre-training resume is out of reach | |
| # either way -- see OPTIMIZER-STATE NOTES at the bottom of this file. | |
| # | |
| # Usage: | |
| # python utils/convert_ckpt_to_nanochat.py \ | |
| # --model checkpoints/model_step005568.pt \ | |
| # --optim checkpoints/optim_step005568.pt \ | |
| # --meta base_checkpoints/d24_decoderstack/meta_005568.json \ | |
| # --out ~/.cache/nanochat/base_checkpoints/d24_decoderstack \ | |
| # --world-size 8 # optional: also emit the optimizer shards | |
| # | |
| # Then, in nanochat (branch fa-varlen): | |
| # from nanochat.checkpoint_manager import build_model | |
| # model, tokenizer, meta = build_model(checkpoint_dir, 5568, device, "eval") | |
| # | |
| # Verify a change to this file with: | |
| # NANOCHAT_PATH=~/nanochat python utils/test_convert_ckpt_to_nanochat.py | |
| # | |
| # The tokenizer is NOT interchangeable with other nanochat d24 releases -- see | |
| # the model card. DecoderStack trained on the 32k vocab shipped with the | |
| # ChrisMcCormick/climbmix_32k_8_170 dataset repo; pairing these weights with a | |
| # different 32k tokenizer produces garbage, not slightly-worse text. | |
| import argparse | |
| import json | |
| import os | |
| import shutil | |
| import torch | |
| def fp32_master(live: torch.Tensor, mantissa: torch.Tensor | None) -> torch.Tensor: | |
| """Rebuild the fp32 master from bf16 live bits + the stashed lower 16 bits. | |
| Mirrors fp32_master() in the training script. int32 rather than uint32 | |
| because CUDA has no uint32 shifts as of torch 2.9; the truncating .to(int16) | |
| and the <<16 discard of sign-extension bits make the two equivalent. With no | |
| mantissa this degrades to a plain upcast (the mantissa bits read as zero), | |
| which is exactly the bf16-precision fallback documented above. | |
| """ | |
| if mantissa is None: | |
| return live.float() | |
| assert mantissa.shape == live.shape, f"mantissa {tuple(mantissa.shape)} != live {tuple(live.shape)}" | |
| bits = (live.view(torch.int16).to(torch.int32) << 16) | \ | |
| (mantissa.view(torch.int16).to(torch.int32) & 0xFFFF) | |
| return bits.view(torch.float32) | |
| def _adamw_groups(ve_slots: int): | |
| """The AdamW half of setup_optimizer()'s group list, in its exact order. | |
| Each entry is (lr_key, [(bank_name, bank_slot), ...], betas, eps, weight_decay). | |
| bank_slot is None for a whole tensor, an int to index a bank's dim 0. The | |
| betas/eps/wd here are constants in BOTH codebases -- nanochat hardcodes them | |
| in setup_optimizer, DecoderStack passes the same numbers to build_schedules -- | |
| so they are not a guess about this run, they are the shared values. | |
| """ | |
| return [ | |
| ("lm_head", [("lm_head", None)], (0.8, 0.96), 1e-10, 0.01), | |
| ("embedding", [("input_embeds", None)], (0.8, 0.995), 1e-10, 0.001), | |
| ("value_embeds", [("value_embeds", j) for j in range(ve_slots)], | |
| (0.8, 0.995), 1e-10, 0.01), | |
| ("resid", [("resid_lambdas", None)], (0.8, 0.95), 1e-10, 0.05), | |
| ("x0", [("x0_lambdas", None)], (0.96, 0.95), 1e-10, 0.0), | |
| ("smear", [("smear_gate", None), ("smear_lambda", None), | |
| ("backout_lambda", None)], (0.8, 0.95), 1e-10, 0.0), | |
| ] | |
| def _matrix_params(n_layer: int, ve: list[int]): | |
| """setup_optimizer()'s `matrix_params`, in list(transformer.h.parameters()) order. | |
| Module registration order gives, per block: attn.c_q, c_k, c_v, c_proj, | |
| [ve_gate], then mlp.c_fc, mlp.c_proj. | |
| """ | |
| out = [] | |
| for i in range(n_layer): | |
| out += [("W_Q", i), ("W_K", i), ("W_V", i), ("W_O", i)] | |
| if i in ve: | |
| out.append(("ve_gate", ve.index(i))) # ve_gate banks by SLOT, not layer | |
| out += [("W_in", i), ("W_out", i)] | |
| return out | |
| def ve_layers(n_layer: int) -> list[int]: | |
| """Layers carrying a value embedding, in bank-slot order. | |
| nanochat's has_ve(): alternating layers, last layer always included. Identical | |
| to StackConfig.ve_layers, and the ascending order matches the VE bank's slot | |
| order (cfg.ve_index), so slot j belongs to layer ve_layers(n_layer)[j]. | |
| """ | |
| return [i for i in range(n_layer) if i % 2 == (n_layer - 1) % 2] | |
| def convert(model_data: dict, optim_state: dict | None) -> dict: | |
| """DecoderStack weights dict -> nanochat state_dict.""" | |
| w = model_data["weights"] | |
| mant = {} if optim_state is None else optim_state | |
| def master(name, i=None): | |
| """fp32 param: live + mantissa. The mantissa of a sharded weight was | |
| all-gathered to full size at capture, so it already lines up 1:1. Pass | |
| `i` to rebuild one slice of a bank -- worth it on the MLP banks, where | |
| materializing the whole thing in fp32 would cost ~0.9 GB per bank.""" | |
| live, m = w[name], mant.get(f"{name}.mantissa") | |
| if i is not None: | |
| live, m = live[i], (None if m is None else m[i]) | |
| return fp32_master(live, m) | |
| n_layer = w["W_Q"].shape[0] | |
| sd = {} | |
| # --- Embeddings: bf16 in nanochat, so the live weights go in untouched. --- | |
| sd["transformer.wte.weight"] = w["input_embeds"] | |
| for slot, layer in enumerate(ve_layers(n_layer)): | |
| sd[f"value_embeds.{layer}.weight"] = w["value_embeds"][slot] | |
| # --- lm_head: AdamW in both, but nanochat keeps it fp32 (it is a Linear, | |
| # not an Embedding, so init_weights() never casts it). --- | |
| sd["lm_head.weight"] = master("lm_head") | |
| # --- Per-layer matrices: unbank dim 0. Every slice is already (out, in). --- | |
| banks = { | |
| "W_Q": "transformer.h.{i}.attn.c_q.weight", | |
| "W_K": "transformer.h.{i}.attn.c_k.weight", | |
| "W_V": "transformer.h.{i}.attn.c_v.weight", | |
| "W_O": "transformer.h.{i}.attn.c_proj.weight", | |
| "W_in": "transformer.h.{i}.mlp.c_fc.weight", | |
| "W_out": "transformer.h.{i}.mlp.c_proj.weight", | |
| } | |
| for name, template in banks.items(): | |
| for i in range(n_layer): | |
| sd[template.format(i=i)] = master(name, i).contiguous() | |
| # --- VE gates: banked by VE SLOT, not by layer, so they unbank through the | |
| # same slot->layer map as the value embeddings. --- | |
| for slot, layer in enumerate(ve_layers(n_layer)): | |
| sd[f"transformer.h.{layer}.attn.ve_gate.weight"] = master("ve_gate", slot).contiguous() | |
| # --- Scalars: fp32-live in both, no mantissa, same names. --- | |
| sd["resid_lambdas"] = w["resid_lambdas"] | |
| sd["x0_lambdas"] = w["x0_lambdas"] | |
| sd["smear_gate.weight"] = w["smear_gate"] | |
| sd["smear_lambda"] = w["smear_lambda"] | |
| sd["backout_lambda"] = w["backout_lambda"] | |
| return sd | |
| def convert_optimizer(model_data: dict, optim_data: dict, world_size: int, rank: int, | |
| lrs: dict) -> dict: | |
| """DecoderStack optimizer capture -> one rank's nanochat optimizer state_dict. | |
| nanochat's state_dict is keyed by flattened param INDEX over setup_optimizer()'s | |
| groups, and DistMuonAdamW shards that state per rank. Our capture all-gathered | |
| everything to full size, so this is re-slicing, not reconstruction. | |
| Returns the dict to torch.save as optim_NNNNNN_rank{rank}.pt. Call once per rank | |
| rather than building them all: at d24/world=8 each shard is ~1 GB. | |
| """ | |
| w, st = model_data["weights"], optim_data["state"] | |
| t_step = optim_data["t_step"] | |
| n_layer = w["W_Q"].shape[0] | |
| ve = ve_layers(n_layer) | |
| n_embd = w["input_embeds"].shape[1] | |
| d_scale = (n_embd / 768) ** -0.5 # setup_optimizer's 1/sqrt(dmodel) AdamW LR scale | |
| def bank(name, slot, attr): | |
| """One param's full-size optimizer state. value_embeds is the odd one out: | |
| its AdamW state is shaped over the FLATTENED (slot * vocab) row axis, so it | |
| has to be folded back to 3-D before a slot can be indexed.""" | |
| t = st[f"{name}.{attr}"] | |
| if slot is None: | |
| return t | |
| if name == "value_embeds": | |
| return t.view(len(ve), -1, t.shape[-1])[slot] | |
| return t[slot] | |
| adamw_lr = { | |
| "lm_head": lrs["unembedding_lr"] * d_scale, | |
| "embedding": lrs["embedding_lr"] * d_scale, | |
| "value_embeds": lrs["embedding_lr"] * d_scale * 0.5, | |
| "resid": lrs["scalar_lr"] * 0.01, | |
| "x0": lrs["scalar_lr"], | |
| "smear": 0.2, # hardcoded in setup_optimizer, not scaled | |
| } | |
| # --- Build the group plan exactly as setup_optimizer() would: AdamW groups in | |
| # a fixed order, then Muon groups keyed by `sorted({p.shape})`. --- | |
| plan = [] # (kind, [(name, slot), ...], hyperparams dict) | |
| for lr_key, params, betas, eps, wd in _adamw_groups(len(ve)): | |
| plan.append(("adamw", params, dict(kind="adamw", lr=adamw_lr[lr_key], | |
| betas=list(betas), eps=eps, weight_decay=wd))) | |
| matrix = _matrix_params(n_layer, ve) | |
| shape_of = lambda p: tuple(w[p[0]].shape[1:]) | |
| for shape in sorted({shape_of(p) for p in matrix}): | |
| plan.append(("muon", [p for p in matrix if shape_of(p) == shape], | |
| dict(kind="muon", lr=lrs["matrix_lr"], momentum=0.95, ns_steps=5, | |
| beta2=0.9, weight_decay=lrs["weight_decay"]))) | |
| # Param indices are assigned by walking the groups in order. | |
| index, idx = {}, 0 | |
| for _, params, _ in plan: | |
| for p in params: | |
| index[p] = idx | |
| idx += 1 | |
| assert idx == 7 + len(ve) + len(matrix), f"param count {idx} does not add up" | |
| state, groups = {}, [] | |
| for kind, params, hp in plan: | |
| groups.append({**hp, "initial_lr": hp["lr"], | |
| "params": [index[p] for p in params]}) | |
| if kind == "adamw": | |
| for p in params: | |
| exp_avg = bank(*p, "exp_avg") | |
| # ZeRO-2: params with >= 1024 elements are row-sharded over dim 0 by | |
| # rank; smaller ones are replicated (nanochat batches those into an | |
| # all_reduce instead of a reduce_scatter). | |
| if exp_avg.numel() >= 1024: | |
| assert exp_avg.shape[0] % world_size == 0, \ | |
| f"{p}: dim 0 ({exp_avg.shape[0]}) must divide world_size {world_size}" | |
| rows = exp_avg.shape[0] // world_size | |
| cut = lambda t, n=rows: t[rank * n:(rank + 1) * n].clone() | |
| else: | |
| cut = lambda t: t.clone() | |
| state[index[p]] = { | |
| "step": t_step, | |
| "exp_avg": cut(exp_avg), | |
| "exp_avg_sq": cut(bank(*p, "exp_avg_sq")), | |
| } | |
| else: | |
| # Muon state is one stacked buffer per GROUP, held under the first | |
| # param's entry, chunked across ranks and zero-padded when the group | |
| # does not divide evenly. | |
| shape = shape_of(params[0]) | |
| chunk = -(-len(params) // world_size) | |
| start = rank * chunk | |
| owned = min(chunk, max(0, len(params) - start)) | |
| mom = torch.zeros(chunk, *shape, dtype=torch.float32) | |
| # nanochat factors the second moment along whichever axis its shape | |
| # heuristic calls the neuron axis; ours is set explicitly per bank. | |
| nc_shape = (shape[-2], 1) if shape[-2] >= shape[-1] else (1, shape[-1]) | |
| snd = torch.zeros(chunk, *nc_shape, dtype=torch.float32) | |
| for k in range(owned): | |
| p = params[start + k] | |
| mom[k] = bank(*p, "frst_mntm") | |
| ours = bank(*p, "scnd_mntm") | |
| if tuple(ours.shape) == nc_shape: | |
| snd[k] = ours | |
| else: | |
| # Only W_O lands here, and only because nanochat's shape | |
| # heuristic picks the other axis on a SQUARE c_proj. That is | |
| # benign: polar express returns a ~orthonormal update, whose | |
| # neuron norms are ~uniform along either axis, so their mean is | |
| # the right common value. On a non-square bank the two axes | |
| # would carry genuinely different information -- refuse. | |
| assert shape[-2] == shape[-1], ( | |
| f"{p}: second-moment axis differs on a non-square bank " | |
| f"{shape} (ours {tuple(ours.shape)}, nanochat {nc_shape}); " | |
| "no faithful conversion exists") | |
| snd[k] = ours.mean() | |
| state[index[params[0]]] = {"momentum_buffer": mom, | |
| "second_momentum_buffer": snd} | |
| return {"state": state, "param_groups": groups} | |
| def main(): | |
| p = argparse.ArgumentParser(description=__doc__) | |
| p.add_argument("--model", required=True, help="DecoderStack model_stepNNNNNN.pt") | |
| p.add_argument("--optim", default=None, | |
| help="matching optim_stepNNNNNN.pt; supplies the mantissas that make " | |
| "the fp32 params exact. Omit for a bf16-precision conversion.") | |
| p.add_argument("--out", required=True, help="output checkpoint dir (nanochat model_tag dir)") | |
| p.add_argument("--meta", default=None, help="meta_NNNNNN.json to copy alongside the model") | |
| p.add_argument("--dump-code", action="store_true", | |
| help="also write the training script embedded in the capture's `code` field") | |
| p.add_argument("--world-size", type=int, default=0, metavar="N", | |
| help="also write optim_NNNNNN_rank{0..N-1}.pt for an N-GPU run " | |
| "(requires --optim). Omit to convert weights only.") | |
| # Group hyperparameters for the emitted optimizer. torch's load_state_dict | |
| # REPLACES param_group dicts with the saved ones, so whatever goes here becomes | |
| # the optimizer's policy on load. Defaults are setup_optimizer()'s own, with | |
| # weight_decay=0.0 -- both the SFT setting and where DecoderStack's cosine-to- | |
| # zero Muon decay actually lands (4.8e-9 at step 5568). nanochat's chat_sft | |
| # restores its own lr right after loading and schedules momentum per step, so | |
| # in practice only betas/eps/weight_decay/ns_steps come from here. | |
| p.add_argument("--unembedding-lr", type=float, default=0.004) | |
| p.add_argument("--embedding-lr", type=float, default=0.2) | |
| p.add_argument("--matrix-lr", type=float, default=0.02) | |
| p.add_argument("--scalar-lr", type=float, default=0.5) | |
| p.add_argument("--weight-decay", type=float, default=0.0) | |
| args = p.parse_args() | |
| if args.world_size and not args.optim: | |
| p.error("--world-size needs --optim (the optimizer state lives in that file)") | |
| # mmap so a 2.8 GB model / 11 GB optimizer file is paged, not slurped. | |
| model_data = torch.load(args.model, map_location="cpu", mmap=True, weights_only=True) | |
| step = model_data["step"] | |
| print(f"loaded {args.model}: step {step}, {len(model_data['weights'])} weights") | |
| optim_data = optim_state = None | |
| if args.optim: | |
| optim_data = torch.load(args.optim, map_location="cpu", mmap=True, weights_only=True) | |
| assert optim_data["step"] == step, f"optim step {optim_data['step']} != model step {step}" | |
| optim_state = optim_data["state"] | |
| n_mant = sum(1 for k in optim_state if k.endswith(".mantissa")) | |
| print(f"loaded {args.optim}: {len(optim_state)} state tensors, {n_mant} mantissas") | |
| else: | |
| print("no --optim: fp32 params will carry bf16 precision (upcast, not exact masters)") | |
| sd = convert(model_data, optim_state) | |
| total = sum(t.numel() for t in sd.values()) | |
| by_dtype = {} | |
| for t in sd.values(): | |
| by_dtype[t.dtype] = by_dtype.get(t.dtype, 0) + t.numel() | |
| print(f"converted: {len(sd)} tensors, {total:,} params " | |
| + ", ".join(f"{n:,} {str(d).replace('torch.', '')}" for d, n in by_dtype.items())) | |
| os.makedirs(args.out, exist_ok=True) | |
| model_path = os.path.join(args.out, f"model_{step:06d}.pt") | |
| torch.save(sd, model_path) | |
| print(f"wrote {model_path} ({os.path.getsize(model_path):,} bytes)") | |
| if args.meta: | |
| meta_path = os.path.join(args.out, f"meta_{step:06d}.json") | |
| shutil.copyfile(args.meta, meta_path) | |
| with open(meta_path, encoding="utf-8") as f: | |
| meta_step = json.load(f)["step"] | |
| assert meta_step == step, f"meta step {meta_step} != model step {step}" | |
| print(f"wrote {meta_path}") | |
| else: | |
| print(f"NOTE: nanochat also needs meta_{step:06d}.json in {args.out} " | |
| "(model_config lives there, not in the .pt)") | |
| if args.world_size: | |
| lrs = dict(unembedding_lr=args.unembedding_lr, embedding_lr=args.embedding_lr, | |
| matrix_lr=args.matrix_lr, scalar_lr=args.scalar_lr, | |
| weight_decay=args.weight_decay) | |
| # One rank at a time -- holding all of them would cost the whole optimizer. | |
| for r in range(args.world_size): | |
| shard = convert_optimizer(model_data, optim_data, args.world_size, r, lrs) | |
| path = os.path.join(args.out, f"optim_{step:06d}_rank{r:d}.pt") | |
| torch.save(shard, path) | |
| print(f"wrote {path} ({os.path.getsize(path):,} bytes)") | |
| if r == 0: | |
| g = shard["param_groups"] | |
| print(f" {len(g)} groups " | |
| f"({sum(1 for x in g if x['kind'] == 'adamw')} adamw / " | |
| f"{sum(1 for x in g if x['kind'] == 'muon')} muon), " | |
| f"{sum(len(x['params']) for x in g)} params, " | |
| f"step {optim_data['t_step']}") | |
| del shard | |
| if args.dump_code: | |
| code_path = os.path.join(args.out, f"code_{step:06d}.py") | |
| with open(code_path, "w", encoding="utf-8", newline="\n") as f: | |
| f.write(model_data["code"]) | |
| print(f"wrote {code_path} (the exact training script for this capture)") | |
| if __name__ == "__main__": | |
| main() | |
| # ----------------------------------------------------------------------------- | |
| # OPTIMIZER-STATE NOTES | |
| # ----------------------------------------------------------------------------- | |
| # SFT CONTINUATION WORKS WITH THE MODEL ALONE. nanochat's chat_sft builds a fresh | |
| # optimizer via model.setup_optimizer() and only optionally warm-starts it from | |
| # load_optimizer_state(); when the shard is absent it prints "optimizer | |
| # checkpoint not found, starting with fresh optimizer (slightly worse)" and | |
| # carries on. --world-size exists to remove that "slightly worse", not to unlock | |
| # anything. | |
| # | |
| # PRE-TRAINING RESUME is genuinely out of reach, and it is the dataloader that | |
| # closes the door, not the optimizer: nanochat's resume needs | |
| # meta_data["dataloader_state_dict"] to put its tokenizing loader back in the | |
| # stream, and DecoderStack reads pre-tokenized binary shards through a loader | |
| # with no equivalent state to hand over. The data order could not be continued no | |
| # matter what the optimizer held. | |
| # | |
| # HYPERPARAMETERS ARE POLICY, NOT STATE. torch's Optimizer.load_state_dict | |
| # REPLACES each param_group dict with the saved one, keeping only 'params' -- so | |
| # whatever this script writes becomes the optimizer's lr/betas/wd on load. That | |
| # is why chat_sft saves and restores its own LRs around the call. The emitted | |
| # groups use setup_optimizer()'s defaults (overridable on the command line), with | |
| # weight_decay=0.0: both the SFT setting and where DecoderStack's cosine-to-zero | |
| # Muon decay actually lands (4.8e-9 at step 5568). betas/eps/adamw-wd are not a | |
| # guess -- they are identical constants in both codebases. | |
| # | |
| # Every buffer we keep has a nanochat counterpart, and the precisions line up on | |
| # everything except the two embedding tables: | |
| # | |
| # DecoderStack nanochat (MuonAdamW / DistMuonAdamW) precision | |
| # ------------------------ -------------------------------------- --------- | |
| # .frst_mntm fp32 Muon state["momentum_buffer"] fp32 both | |
| # .scnd_mntm fp32 Muon state["second_momentum_buffer"] fp32 both | |
| # .exp_avg fp32 AdamW state["exp_avg"] see below | |
| # .exp_avg_sq fp32 AdamW state["exp_avg_sq"] see below | |
| # .mantissa uint16 (no counterpart -- nanochat's fp32 param IS | |
| # the master; consumed above to rebuild it) | |
| # | |
| # nanochat allocates its Muon buffers as `dtype=p.dtype` and its AdamW buffers as | |
| # `torch.zeros_like(p)`. Its Muon params and lm_head are fp32, so those match us. | |
| # But wte and value_embeds are bf16 PARAMS, so THEIR AdamW moments are bf16 -- | |
| # where ours are fp32. That is the one precision difference, and ours is the more | |
| # precise of the two, deliberately: it is only the GRADIENTS that are bf16 for | |
| # those two tables (they are the biggest tensors in the model, so fp32 grads | |
| # would double their scatter and comm traffic, and bf16 matches the autograd | |
| # baseline's numerics). The moment math stays fp32 -- adamw_step_fused upcasts on | |
| # the way in, `grad = grad.to(exp_avg.dtype)`. There is no bf16 AdamW variant in | |
| # the file: the two AdamW kernels differ in whether the param carries a mantissa | |
| # (adamw_step_fused vs adamw_step_fused_fp32), not in moment dtype. | |
| # | |
| # W_O's reduction axis differs between the two, and at d24 it costs nothing. | |
| # NorMuon's factored second moment is a per-neuron mean-square; nanochat infers | |
| # the neuron axis from the shape (`red_dim = -1 if shape[-2] >= shape[-1] else | |
| # -2`) while DecoderStack states it (m.W_O.residual_dim = -2), so the two | |
| # disagree on a square c_proj -- ours is (1, 1536) where nanochat's is (1536, 1). | |
| # But polar express returns a ~orthonormal update, and a square orthonormal | |
| # matrix has ~uniform neuron norms along either axis: there is no variance to | |
| # reduce, the rescale is a ~no-op, and the run is unaffected by the choice. | |
| # The explicit axis earns its keep only when n_heads * d_head != d_model. Above | |
| # d_model, W_O looks like an MLP projection and the shape heuristic happens to | |
| # agree; below it the heuristic picks the wrong axis, and since W_O stores its | |
| # heads transposed relative to QKV, the right answer is not one a shape alone | |
| # can give. Every other bank agrees at d24 (W_in -1, W_out -2, QKV -1). | |
| # | |
| # ve_gate is the second place the axes can diverge, and it is worth knowing about | |
| # because it is NOT square in general. DecoderStack banks it (num_ves, n_kv_heads, | |
| # d_ve_gate) with residual_dim = -1, so the neurons are the n_kv_heads rows; | |
| # nanochat's Linear(ve_gate_channels=12, n_kv_head) hits the same heuristic and | |
| # agrees only when n_kv_head >= 12. At d24 n_kv_head == 12 == d_ve_gate, so the | |
| # bank is square and the two land together. A model with fewer than 12 KV heads | |
| # would disagree for real -- convert_optimizer() asserts rather than papering | |
| # over it, since outside the square case the two axes carry different | |
| # information. | |
| # | |
| # --world-size implements the mapping below. It is mechanical but fiddly, because | |
| # nanochat's state_dict is keyed by flattened param INDEX and is sharded per rank: | |
| # - Param order is setup_optimizer()'s group order: lm_head, wte, | |
| # value_embeds.*, resid_lambdas, x0_lambdas, [smear_gate.weight, | |
| # smear_lambda, backout_lambda], then the Muon groups in `sorted({shapes})` | |
| # order -- (12,12) ve_gates, (1536,1536) c_q/c_k/c_v/c_proj interleaved in | |
| # block order, (1536,6144) mlp.c_proj, (6144,1536) mlp.c_fc. | |
| # - AdamW state for params with >= 1024 elements is sliced over dim 0 by rank; | |
| # smaller ones are replicated. Ours is captured all-gathered to full size, | |
| # so it just needs re-slicing (value_embeds first reshaped from its flattened | |
| # (num_ves * vocab, kv_dim) row axis back to (num_ves, vocab, kv_dim)). | |
| # - Muon state is stacked per group and chunked: rank r owns params | |
| # [r*ceil(K/W) : (r+1)*ceil(K/W)] of the group, zero-padded when K % W != 0. | |