Declarative checkpoint config conversion

[jlamypoirier]

Summary

Replaces the imperative import_config / export_config methods on checkpoint converter classes with declarative ConfigConverter primitives keyed off FieldHint.architecture. Two recursive coverage walkers — one over Fast-LLM architecture fields, one over the HF dict — guarantee that every architecture-significant field and every HF dict key is accounted for on import and export. Unsupported feature detection is driven from a single source instead of scattered per-converter asserts.

What's in

• Primitives (fast_llm/engine/checkpoint/external.py): Rename, ConstantImport, ConstantExport, Optional, Default, Ignored, ImportOnly, Custom, plus the recursive Nested / Dispatch / TypedDictContainer. Each declares the Fast-LLM paths and HF paths it consumes so the walkers can verify coverage on both sides.
• ConfigSectionConverter ABC: cached _create_config_converters per class (subclasses extend the parent's dict by re-declaring keys), _validate_export hook for format-specific cross-field invariants, and recursive check_architecture_coverage / check_hf_coverage.
• Architecture-hint reclassification (eight fields): attention dense_layer and softmax_scale_power, MLP activation, MoE router, Llama3 / Yarn rotary scaling parameters, StochasticMixer.main_mixer_name, vision patch_height / patch_width, and PatternBlockSequenceConfig.blocks. These are what the coverage walkers now require to be consumed.
• GPT-side migration: llama, mistral, qwen2 (including the MRoPE guard declaratively as an ImportOnlyConfigConverter), mtp_llama, mixtral, diffusion_dream, diffusion_llama, apriel2 (text), the apriel hybrid-SSM mixers (Mamba, GatedDeltaNet, KimiDeltaAttention), and gemma4 — the latter as a declarative shell over imperative cross-block helpers (Gemma 4's HF format merges sliding + full attention into one config and cross-references hidden_size).
• Multimodal migration: Llava (vision adapter, vision model, base model) and multimodal Apriel2 (vision attention/block/MLP/encoder/embeddings/adapter/model and the top-level multimodal base) are now ConfigSectionConverters. The Llava adapter is declared at the base-model scope so its intermediate_size import can cross-reference text_config.hidden_size. Apriel2 vision injects patch_size / max_image_size rotary metadata at the vision-model scope (the smallest scope that sees both embeddings.patch_height and the encoder rotary subtree).
• Decoder dispatcher flattening: Mistral / Qwen2 / Mixtral DecoderConverter subclasses deleted. LlamaBaseModelConverter inlines the Fixed/Pattern dispatch parameterised by block_converter_class. The imperative LlamaDecoderConverter helper class survives because Pixtral's vision encoder dispatch and Apriel's per-position hybrid_block_layout dispatch don't fit the inlined common case; AprielBaseModelConverter overrides the "decoder" declaration to delegate to AprielDecoderConverter via a new apriel_decoder_converter_class ClassVar.
• HF coverage allowlist propagation: check_hf_coverage now matches the allowlist against any segment of the walked path (not just top-level keys). transformers' generic PretrainedConfig metadata (architectures, torch_dtype, transformers_version, …) is now accepted at any depth — relevant for Llava's vision_config and similar nested sub-configs that transformers auto-populates on round-trip.
• Static test (tests/models/test_converters.py): walks every registered handler's converter tree, runs check_architecture_coverage on each section node, and validates that every OptionalConfigConverter sentinel matches the field's resolved default (so an exported sentinel-equal value can't silently drift on re-import).

Main merge

• Brings in the new architecture fields landed since branch cut: AttentionConfig.{query,key,value}_norm + shared_key_value, MLPConfig.{pre,post}_norm, DecoderBlockConfig.{pre,post}_{mixer,mlp}_normalization + output_scale, LanguageModelEmbeddingsConfig.embedding_scale, LanguageModelHeadConfig.final_logit_softcap, MoEMLPConfig.{router_normalization, router_scale, router_input_scale, router_per_expert_scale}. Each new field is claimed on every section converter whose fast_llm_config_class is the field's container via ConstantImportConfigConverter (asserts the Fast-LLM default on export, injects the default on import); OptionalParameterConfig fields (output_scale, router_scale, router_per_expert_scale) are claimed via IgnoredConfigConverter + a _validate_export not enabled assertion to preserve the imperative reject-non-default behaviour.

Notable shape decisions

• Coverage check is type-strict (type(config) is cls.fast_llm_config_class). Strict subclasses defer to a more specific converter, which kept partially migrated callers working through super() during the multi-step migration.
• HF-side coverage walker consumes prefixes registered by every declaration, descending through Nested / Dispatch sub-converters under their hf_path. IgnoredConfigConverter(hf_paths=...) is the explicit opt-out for HF-only fields with no Fast-LLM counterpart (Mixtral router-runtime toggles, Qwen2 sliding-window, …) and for the model-specific defaults transformers' Pixtral/Llava fill on round-trip (head_dim, image_size, layer_norm_eps, projection_dim, vocab_size, image_seq_length, tie_word_embeddings at the Llava level).
• _create_config_converters is @functools.cached. Subclasses must return a fresh dict ({**super()._create_config_converters(), ...}) — mutating the parent's returned dict would corrupt its cache entry. Documented on the base method.
• IgnoredConfigConverter is recursive (recurses=True). Used for sub-configs with no architecture leaves (ParameterConfig, Mixtral's router sub-config) and for HF-only fields. Non-architecture fields (lr_scale, apply_peft, initialization sub-config) are by design not part of the HF round-trip; Fast-LLM keeps them on the in-memory config independently.

What's deferred

• Weight-converter declarative refactor. The weight side still uses today's WeightConverter subclasses (SplitWeightConverter, MLPLayer2Converter, KeyValueWeightConverter, …) and the per-converter get_converters classmethods.
• Gemma 4 / Apriel block-level declarative. Gemma 4's Gemma4BlockConverter / Gemma4DecoderConverter and Apriel's AprielBlockConverter / AprielDecoderConverter remain imperative — Gemma 4 because the HF format merges sliding + full attention into one config and cross-references hidden_size (doesn't fit per-section decomposition), Apriel because hybrid_block_layout is a positional list discriminator (would need a new ListDispatchConfigConverter primitive).
• A handful of structural follow-ups recorded for future review rounds: IgnoredConfigConverter as a _ignored_fields ClassVar, get_converters signature uniformity, LlamaDecoderConverter as a ConfigSectionConverter, IgnoredConfigConverter default-round-trip maintenance test.

Test plan

• pytest -v -n 8 tests/models/test_checkpoint.py tests/models/test_hf_roundtrip.py tests/models/test_converters.py: 271 passed, no failures (gemma4 dependency group now included).
• pytest -v -n 8 fast_llm_external_models/tests/: 2109 passed, 42 skipped (separate invocation per CLAUDE.md).
• Manual smoke: fast-llm convert --input.format <fmt> --input.path <ref> --output.format <fmt> --output.path <tmp>; reload both and compare configs.

🤖 Generated with Claude Code