karpathy/nanochat · error · ValueError

Unknown optimizer kind

Error message

Unknown optimizer kind: {group['kind']}

What it means

The fused MuonAdamW optimizer in nanochat/optim.py partitions parameters into groups, each tagged with group['kind'] of 'adamw' or 'muon', and dispatches per-kind reduce operations (all-reduce of grads/optimizer state) in `step()`. A param group whose 'kind' key is missing or misspelled hits the else and raises ValueError during the reduce phase (Phase 1).

Solutions

  1. Ensure every param group dict passed to MuonAdamW has kind='adamw' or kind='muon' (lowercase, exact).
  2. Use the canonical group-construction code in nanochat/gpt.py (`configure_optimizers`) as the template.
  3. Add an assertion over param_groups before training: all(g['kind'] in ('adamw','muon') for g in opt.param_groups).

Example fix

# before
opt = MuonAdamW([dict(params=lm_head_params, lr=6e-4)])

# after
opt = MuonAdamW([dict(kind='adamw', params=lm_head_params, lr=6e-4, betas=(0.8, 0.96), eps=1e-10, weight_decay=0.01)])
Defensive patterns

Strategy: validation

Validate before calling

VALID_KINDS = ('adamw', 'muon')
for g in param_groups:
    assert g.get('kind') in VALID_KINDS, f"param group kind must be one of {VALID_KINDS}, got {g.get('kind')!r}"
opt = MuonAdamW(param_groups)

Type guard

def is_valid_param_group(group: dict) -> bool:
    return isinstance(group, dict) and group.get('kind') in ('adamw', 'muon')

Prevention

When it happens

Trigger: Constructing MuonAdamW with a param group dict lacking a 'kind' key (KeyError becomes None via group['kind'] only if key exists — actually a missing key raises KeyError first; the ValueError fires on any wrong value such as 'sgd', 'AdamW' with wrong casing, or 'lion').

Common situations: Writing custom training scripts that build param groups for MuonAdamW by hand; casing or spelling mistakes in 'kind'; copying param groups from torch.optim.AdamW-style code where 'kind' does not exist (that yields KeyError, but renamed/edited dicts yield this ValueError).

Related errors

AI-assisted analysis of karpathy/nanochat@92d63d4e8b (2026-08-15). Data as JSON: /api/errors/64a628098bd4c97a. Report an issue: GitHub.

Appendix: source

Thrown at nanochat/optim.py:446

    @torch.no_grad()
    def step(self):
        # On a single rank (no multi-rank process group), all communication is skipped
        if dist.is_available() and dist.is_initialized():
            rank = dist.get_rank()
            world_size = dist.get_world_size()
        else:
            rank = 0
            world_size = 1

        # Phase 1: launch all async reduce ops
        reduce_infos: list[dict] = []
        for group in self.param_groups:
            if group['kind'] == 'adamw':
                reduce_infos.append(self._reduce_adamw(group, world_size))
            elif group['kind'] == 'muon':
                reduce_infos.append(self._reduce_muon(group, world_size))
            else:
                raise ValueError(f"Unknown optimizer kind: {group['kind']}")

        # Phase 2: wait for reduces, compute updates, launch gathers
        gather_list: list[dict] = []
        for group, info in zip(self.param_groups, reduce_infos):
            if group['kind'] == 'adamw':
                self._compute_adamw(group, info, gather_list, rank, world_size)
            elif group['kind'] == 'muon':
                self._compute_muon(group, info, gather_list, rank)
            else:
                raise ValueError(f"Unknown optimizer kind: {group['kind']}")

        # Phase 3: wait for gathers, copy back
        self._finish_gathers(gather_list)

View on GitHub (pinned to 92d63d4e8b)