bridge.recipes.kimi.kimi_k2

Module Contents

Functions

_get_kimi_k2_pipeline_layout	Get pipeline layout for Kimi-K2 based on PP and VP size.
kimi_k2_pretrain_config	Return a pre-training config for Kimi-K2 (1T).

API

bridge.recipes.kimi.kimi_k2._get_kimi_k2_pipeline_layout(pp_size: int, vp_size: int)

Get pipeline layout for Kimi-K2 based on PP and VP size.

bridge.recipes.kimi.kimi_k2.kimi_k2_pretrain_config(

optimizer_type: str = 'muon',

) → megatron.bridge.training.config.ConfigContainer

Return a pre-training config for Kimi-K2 (1T).

Recommended parallelism: TP=2, PP=16, EP=32

Parameters:

optimizer_type – ‘adam’ or ‘muon’ (default).