Full-State Safetensor → MaxText Checkpoint Converter

背景与动机

在 Ling3 模型训练中，我们需要从 Megatron DCP 框架迁移到 MaxText 进行训练。此前已有 model parameters 的转换能力（to_maxtext.py），但无法迁移 optimizer state（Muon momentum、AdamW 一/二阶矩），导致每次框架切换后模型必须从头 warm-up optimizer，产生训练 spike 和算力浪费。

本设计实现了完整训练状态（model params + optimizer state + step counters）的无损迁移。

转换管线

text

Megatron DCP
    │
    ▼  convert_dcp_to_full_state_safetensor.py (上游)
ling3-muon-full-state-v1.safetensors
    │
    ▼  to_maxtext_ling3_muon_full_state.py (本 PR)
MaxText Orbax full-state checkpoint
    │
    ▼  load_full_state_path
Resume training (step counter + optimizer history preserved)

核心设计

1. Source Safetensor Schema

采用带命名空间前缀的 key 组织方式：

前缀	含义
`param/<hf_key>`	模型参数
`optim/muon_mu/<hf_key>`	Muon 动量
`optim/adam_mu/<hf_key>`	AdamW 一阶矩
`optim/adam_nu/<hf_key>`	AdamW 二阶矩
`meta/step`	训练步数
`meta/muon_count`, `meta/adam_count`	优化器计数器

Metadata 字段：format=ling3-muon-full-state-v1, model_name, optimizer, step, consumed_train_samples

2. 架构组件

SourceSafetensorReader

支持单文件和多分片（model.safetensors.index.json）模式
使用 JAX framework 读取 bfloat16 tensor
自动 reshape Megatron DistributedOptimizer 的 flattened 1D buffers
Metadata 校验（format version、model_name、optimizer type）

Ling3FullStateManifest

复用 param_mapping.py 的 HF→MaxText 映射关系
分类参数属于 Muon partition 还是 Adam partition
可生成 JSON manifest 用于调试和验证

SlotTransformer

Single-pass transform：一次遍历同时处理 params + optimizer slots
复用 to_maxtext.py 暴露的 5 个公共 API 函数
相比 4 次独立遍历，减少 GCS FUSE 上的 shard file 打开次数，提升 I/O locality

OptStatePatcher

将转换后的 optimizer tensor 按 semantic path 匹配到 MaxText 的 opt_state 树中
处理 NamedTuple 字段的 dot 前缀（.mu, .nu）
设置 step counter（Muon/Adam 可独立计数）
严格校验：若 0 个 leaf 被 patch 则抛出 CRITICAL 错误

3. Optimizer Partition 分类策略

Ground truth 来自 source safetensors keys，而非 MaxText 侧的 transform_logic。

原因：DCP converter 基于 Megatron 的实际 optimizer 分组产出 optim/muon_mu/ 和 optim/adam_mu/ 前缀，这是最权威的分类来源。通过反查 param_map 将 HF key 映射回 MaxText key，建立每个参数的 partition 归属。

4. `to_maxtext.py` 重构

将 5 个内部函数暴露为 public API，保持 backward-compatible aliases：

build_multi_axis_stacked_tensor
build_single_axis_stacked_tensor
get_hf_loading_function
get_maxtext_indices_and_shapes
get_maxtext_weight

5. 训练恢复支持

新增 consumed_train_samples_override 配置字段
仅在无本地 checkpoint 时生效（首次 full-state 加载）
Job self-healing 时优先使用 restored trainer_state，避免错误重置 sample counter
所有 pretrain/submit 脚本支持 LOAD_FULL_STATE_PATH 环境变量

验证结果

验证项	结果
Round-trip (safetensors → Orbax → HF safetensors)	全部 9686 权重精确匹配
8-chip 训练	loss ~1.4，无 spike，从 step 454212 平滑恢复
64-chip 训练	loss ~1.4，grad norm ~0.07 稳定
1024-chip 生产运行	loss ~1.22，116 TFLOP/s/device，2.67s/step，step 457000+ 稳定运行

使用方式

bash

# 生成 manifest（不执行转换）
python3 -m maxtext.checkpoint_conversion.to_maxtext_ling3_muon_full_state \
  src/MaxText/configs/base.yml \
  model_name=ling3-flash opt_type=muon scan_layers=true \
  base_output_directory=/path/to/out hardware=cpu \
  skip_jax_distributed_system=true attention=dot_product \
  --source_safetensor=/path/to/full_state.safetensors \
  --emit_manifest=/path/to/manifest.json \
  --manifest_only=true

# 完整转换
python3 -m maxtext.checkpoint_conversion.to_maxtext_ling3_muon_full_state \
  src/MaxText/configs/base.yml \
  model_name=ling3-flash opt_type=muon scan_layers=true \
  base_output_directory=gs://bucket/ckpt hardware=cpu \
  skip_jax_distributed_system=true attention=dot_product \
  --source_safetensor=/path/to/full_state.safetensors

# 使用转换后的 checkpoint 恢复训练
LOAD_FULL_STATE_PATH=gs://bucket/ckpt bash scripts/pretrain_ling3_flash.sh

文件清单

文件	说明
`src/maxtext/checkpoint_conversion/to_maxtext_ling3_muon_full_state.py`	核心转换脚本 (1146 行)
`src/maxtext/checkpoint_conversion/to_maxtext.py`	重构为 public API
`src/maxtext/configs/types.py`	新增 `consumed_train_samples_override`
`src/maxtext/trainers/pre_train/train.py`	支持 override 逻辑
`tests/unit/.../ling3_muon_full_state_conversion_test.py`	单元测试 (282 行)
`.github/ci/tpu-*-fullstate-job.yaml`	64/1024 chip 任务模板
`scripts/submit_ling3_*_fullstate_job.sh`	提交脚本
`scripts/pretrain_*.sh`	所有 pretrain 脚本支持 LOAD_FULL_STATE_PATH

设计决策

Single-pass vs Multi-pass：选择 single-pass 以优化 GCS FUSE I/O（减少 4x shard open）
Ground-truth partition from source keys：比依赖 MaxText transform_logic 更可靠，因为源头是 DCP converter 的实际分组
Strict validation：若没有任何 optimizer leaf 被 patch，则 fail-fast，避免静默丢失 optimizer state 导致训练发散
consumed_train_samples_override 仅首次生效：避免 job restart 时错误覆盖已恢复的 trainer_state

Full-State Safetensor → MaxText Checkpoint Converter ​

背景与动机 ​

转换管线 ​

核心设计 ​

1. Source Safetensor Schema ​

2. 架构组件 ​

SourceSafetensorReader ​

Ling3FullStateManifest ​

SlotTransformer ​

OptStatePatcher ​

3. Optimizer Partition 分类策略 ​

4. to_maxtext.py 重构 ​

5. 训练恢复支持 ​

验证结果 ​

使用方式 ​

文件清单 ​

设计决策 ​