[dist] fix: make OptimizerState EP-dim aware to fix its dcp saving (#228)

Similar to ModelState, OptimizerState also needs to be EP-dim aware so that dcp can save it properly.
Now it has same workflow with ModelState:
    * before saving with dcp, restore EP dim
    * after loading state dict from dcp, drop EP dim

Author: Luosuu

.github/workflows/gpu_unit_tests.yml

@@ -77,6 +77,9 @@ jobs:
       - name: Run models tests
         run: |
           pytest -s -x tests/models/test_models_patch.py
+      - name: Run e2e dcp save and load test
+        run: |
+          pytest -s -x tests/checkpoints/test_trainer_saveload.py
 
 
   cleanup:

configs/model_configs/qwen/qwen3_moe_30a3b_4_layers.json

@@ -0,0 +1,38 @@
+{
+  "architectures": [
+    "Qwen3MoeForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "decoder_sparse_step": 1,
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 6144,
+  "max_position_embeddings": 262144,
+  "max_window_layers": 48,
+  "mlp_only_layers": [],
+  "model_type": "qwen3_moe",
+  "moe_intermediate_size": 768,
+  "norm_topk_prob": true,
+  "num_attention_heads": 32,
+  "num_experts": 128,
+  "num_experts_per_tok": 8,
+  "num_hidden_layers": 4,
+  "num_key_value_heads": 4,
+  "output_router_logits": false,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 10000000,
+  "router_aux_loss_coef": 0.001,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

tests/checkpoints/ep4.yaml

@@ -0,0 +1,34 @@
+model:
+  # model_path: ./qwen3moe_4layers_merged
+  config_path: configs/model_configs/qwen/qwen3_moe_30a3b_4_layers.json
+  tokenizer_path: Qwen/Qwen3-30B-A3B
+  weight_path: None
+  moe_implementation: fused
+  attn_implementation: flash_attention_2
+
+data:
+  train_path: dummy
+  max_seq_len: 128
+
+train:
+  output_dir: ./test_trainer_saveload_ep4
+  data_parallel_mode: fsdp2
+  expert_parallel_size: 4
+  enable_full_shard: true
+  init_device: meta
+  global_batch_size: 8
+  micro_batch_size: 1
+  rmpad: false
+  rmpad_with_pos_ids: true
+  dyn_bsz_margin: 0
+  lr: 3.0e-4
+  lr_warmup_ratio: 0.007
+  lr_decay_style: constant
+  lr_decay_ratio: 1.0
+  weight_decay: 0.01
+  max_grad_norm: 1.0
+  use_wandb: false
+  enable_profiling: false
+  max_steps: 5
+  ckpt_manager: dcp
+  save_async: true

tests/checkpoints/ep8.yaml

@@ -0,0 +1,34 @@
+model:
+  # model_path: ./qwen3moe_4layers_merged
+  config_path: configs/model_configs/qwen/qwen3_moe_30a3b_4_layers.json
+  tokenizer_path: Qwen/Qwen3-30B-A3B
+  weight_path: None
+  moe_implementation: fused
+  attn_implementation: flash_attention_2
+
+data:
+  train_path: dummy
+  max_seq_len: 128
+
+train:
+  output_dir: ./test_trainer_saveload_ep8
+  data_parallel_mode: fsdp2
+  expert_parallel_size: 8
+  enable_full_shard: true
+  init_device: meta
+  global_batch_size: 8
+  micro_batch_size: 1
+  rmpad: false
+  rmpad_with_pos_ids: true
+  dyn_bsz_margin: 0
+  lr: 3.0e-4
+  lr_warmup_ratio: 0.007
+  lr_decay_style: constant
+  lr_decay_ratio: 1.0
+  weight_decay: 0.01
+  max_grad_norm: 1.0
+  use_wandb: false
+  enable_profiling: false
+  max_steps: 5
+  ckpt_manager: dcp
+  save_async: true

tests/checkpoints/no_ep.yaml

@@ -0,0 +1,34 @@
+model:
+  # model_path: ./qwen3moe_4layers_merged
+  config_path: configs/model_configs/qwen/qwen3_moe_30a3b_4_layers.json
+  tokenizer_path: Qwen/Qwen3-30B-A3B
+  weight_path: None
+  moe_implementation: fused
+  attn_implementation: flash_attention_2
+
+data:
+  train_path: dummy
+  max_seq_len: 128
+
+train:
+  output_dir: ./test_trainer_saveload_no_ep
+  data_parallel_mode: fsdp2
+  expert_parallel_size: 1
+  enable_full_shard: true
+  init_device: meta
+  global_batch_size: 8
+  micro_batch_size: 1
+  rmpad: false
+  rmpad_with_pos_ids: true
+  dyn_bsz_margin: 0
+  lr: 3.0e-4
+  lr_warmup_ratio: 0.007
+  lr_decay_style: constant
+  lr_decay_ratio: 1.0
+  weight_decay: 0.01
+  max_grad_norm: 1.0
+  use_wandb: false
+  enable_profiling: false
+  max_steps: 5
+  ckpt_manager: dcp
+  save_async: true

tests/utils/test_trainer_saveload.py …sts/checkpoints/test_trainer_saveload.pytests/utils/test_trainer_saveload.py renamed to tests/checkpoints/test_trainer_saveload.py tests/utils/test_trainer_saveload.py renamed to tests/checkpoints/test_trainer_saveload.py

@@ -1,5 +1,6 @@
 import json
 import os
+import subprocess
 from dataclasses import asdict, dataclass, field
 from typing import Any, Dict, List, Optional
 
@@ -8,7 +9,7 @@
 from tqdm import trange
 
 from veomni.checkpoint import build_checkpointer
-from veomni.data import build_dummy_dataset, build_streaming_dataloader
+from veomni.data import build_dataloader, build_dummy_dataset
 from veomni.distributed.offloading import build_activation_offloading_context
 from veomni.distributed.parallel_state import get_parallel_state, init_parallel_state
 from veomni.distributed.torch_parallelize import build_parallelize_model
@@ -21,37 +22,24 @@
 
 
 """
-torchrun --nnodes=1 --nproc-per-node=8 --master-port=4321 tests/utils/test_trainer_saveload.py \
---model.model_path Qwen/Qwen3-4B \
---train.expert_parallel_size 1 \
---train.global_batch_size 8 \
---train.micro_batch_size 1 \
---data.max_seq_len 128 \
---data.train_path "dummy" \
---train.output_dir ./test_trainer_saveload \
---train.max_steps 5 \
---train.rmpad false \
---train.rmpad_with_pos_ids true \
---train.data_parallel_mode "fsdp2" \
---train.init_device "meta" \
---train.ckpt_manager "dcp"
-
-torchrun --nnodes=1 --nproc-per-node=8 --master-port=4321 tests/utils/test_trainer_saveload.py \
---model.model_path /path/to/Qwen3-30B-A3B-Instruct-2507-merge \
+torchrun --nnodes=1 --nproc-per-node=8 --master-port=4321 tests/checkpoints/test_trainer_saveload.py \
+--model.config_path configs/model_configs/qwen/qwen3_moe_30a3b_4_layers.json \
+--model.weight_path None \
+--model.tokenizer_path /mnt/hdfs/models/Qwen3-30B-A3B \
 --model.moe_implementation fused \
 --model.attn_implementation flash_attention_2 \
---train.expert_parallel_size 4 \
+--train.expert_parallel_size 8 \
 --train.global_batch_size 8 \
 --train.micro_batch_size 1 \
 --data.max_seq_len 128 \
 --data.train_path "dummy" \
---train.output_dir ./test_trainer_saveload \
+--train.output_dir ./test_trainer_saveload_ep8 \
 --train.max_steps 5 \
 --train.rmpad false \
 --train.rmpad_with_pos_ids true \
 --train.data_parallel_mode "fsdp2" \
 --train.init_device "meta" \
---train.ckpt_manager "dcp"
+--train.ckpt_manager "dcp" $@ 2>&1 | tee test_saveload_ep8.log
 """
 
 # To prevent DCP from complaining "too many open files"
@@ -143,8 +131,9 @@ def main():
     train_dataset = build_dummy_dataset(task_type="text", size=train_data_size, max_seq_len=args.data.max_seq_len)
 
     args.train.compute_train_steps(args.data.max_seq_len, args.data.train_size)
-    train_dataloader = build_streaming_dataloader(
+    train_dataloader = build_dataloader(
         dataset=train_dataset,
+        dataloader_type="streaming",
         micro_batch_size=args.train.micro_batch_size,
         global_batch_size=args.train.global_batch_size,
         dataloader_batch_size=args.train.dataloader_batch_size,
@@ -356,5 +345,44 @@ def step_id(s):
     dist.destroy_process_group()
 
 
+def test_trainer_saveload_ep8():
+    ep8_command = [
+        "torchrun",
+        "--nnodes=1",
+        "--nproc_per_node=8",
+        "--master_port=4321",
+        "tests/utils/test_trainer_saveload.py",
+        "tests/checkpoints/ep8.yaml",
+    ]
+    ep8_result = subprocess.run(ep8_command, check=True)
+    assert ep8_result.returncode == 0
+
+
+def test_trainer_saveload_ep4():
+    ep4_command = [
+        "torchrun",
+        "--nnodes=1",
+        "--nproc_per_node=8",
+        "--master_port=4321",
+        "tests/checkpoints/test_trainer_saveload.py",
+        "tests/checkpoints/ep4.yaml",
+    ]
+    ep4_result = subprocess.run(ep4_command, check=True)
+    assert ep4_result.returncode == 0
+
+
+def test_trainer_saveload_no_ep():
+    no_ep_command = [
+        "torchrun",
+        "--nnodes=1",
+        "--nproc_per_node=8",
+        "--master_port=4321",
+        "tests/checkpoints/test_trainer_saveload.py",
+        "tests/checkpoints/no_ep.yaml",
+    ]
+    no_ep_result = subprocess.run(no_ep_command, check=True)
+    assert no_ep_result.returncode == 0
+
+
 if __name__ == "__main__":
     main()