bridge.models.kimi.kimi_k3_bridge#

Hugging Face ↔ Megatron conversion for the Kimi K3 language backbone.

The initial parameter mapping was adapted from the Apache-2.0 Miles implementation: https://github.com/radixark/miles/blob/dc62a0bd4b7af1c59ee2084852eb18b5585ec082/miles_plugins/mbridge/kimi_k3.py

Module Contents#

Classes#

KimiK3Bridge

Megatron Bridge for the Kimi K3 language backbone.

API#

class bridge.models.kimi.kimi_k3_bridge.KimiK3Bridge#

Bases: megatron.bridge.models.conversion.model_bridge.MegatronModelBridge

Megatron Bridge for the Kimi K3 language backbone.

_HF_PASSTHROUGH_PREFIXES#

(‘vision_tower.’, ‘mm_projector.’)

classmethod hf_to_megatron_activation(hidden_act: str)#

Use SiLU as the config marker; the K3 spec installs the exact SiTU module.

classmethod megatron_to_hf_activation(activation_func) → str#

Serialize the custom activation with its Hugging Face name.

hf_config_to_provider_kwargs(hf_config) → dict#

Map the nested K3 language configuration with the common config mappings.

provider_bridge(
hf_pretrained: megatron.bridge.models.hf_pretrained.causal_lm.PreTrainedCausalLM,
) → megatron.bridge.models.kimi.kimi_k3_provider.KimiK3ModelProvider#

Translate the nested Kimi K3 text configuration.

mapping_registry() → megatron.bridge.models.conversion.mapping_registry.MegatronMappingRegistry#

Map K3’s nested language model and custom layer parameters.

static _load_one_hf_weight(
name: str,
hf_state_dict: collections.abc.Mapping[str, torch.Tensor],
) → torch.Tensor#
maybe_modify_loaded_hf_weight(
hf_param: str | dict[str, str],
hf_state_dict: collections.abc.Mapping[str, torch.Tensor],
) → torch.Tensor | dict[str, torch.Tensor]#

Load K3 weights, dequantizing routed experts from MXFP4 when needed.

maybe_modify_converted_hf_weight(
task: megatron.bridge.models.conversion.model_bridge.WeightConversionTask,
converted_weights_dict: dict[str, torch.Tensor],
hf_state_dict: collections.abc.Mapping[str, torch.Tensor],
) → dict[str, torch.Tensor]#

Restore KDA padding and the source routed-expert MXFP4 representation.

build_conversion_tasks(
hf_pretrained,
megatron_model,
weight_dtype=None,
) → list[megatron.bridge.models.conversion.model_bridge.WeightConversionTask]#

Expose virtual BF16 expert weights for MXFP4 packed/scale pairs.

stream_weights_megatron_to_hf(
megatron_model: megatron.core.models.gpt.gpt_model.GPTModel | list[megatron.core.models.gpt.gpt_model.GPTModel],
hf_pretrained: megatron.bridge.models.hf_pretrained.causal_lm.PreTrainedCausalLM,
cpu: bool = True,
show_progress: bool = True,
conversion_tasks: list[megatron.bridge.models.conversion.model_bridge.WeightConversionTask] | None = None,
merge_adapter_weights: bool = True,
weight_dtype: torch.dtype | None = None,
) → collections.abc.Iterable[megatron.bridge.models.conversion.model_bridge.HFWeightTuple]#

Export the language model and preserve unchanged multimodal weights.