[Question] How to properly resume LoRA training from checkpoint in Megatron with Qwen3-Omni and dataset packing + streaming

Environment

• Model: Qwen/Qwen3-Omni-30B-A3B-Instruct
• Training Type: LoRA (rank=8, alpha=32, target_modules=all-linear)
• Framework: ms-swift with Megatron backend
• Hardware: 2x GPU setup with expert_model_parallel_size=2
• Swift Version: 3.9.0.dev0
• Megatron Core Version: 0.13.2

Problem Description

I'm unable to properly resume LoRA training from a checkpoint. The training starts from iteration 1 instead of continuing from the saved checkpoint iteration (200), and optimizer/scheduler states appear to be reset.

Current Training Command

PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
MAX_PIXELS=1003520 VIDEO_MAX_PIXELS=50176 FPS_MAX_FRAMES=12 CUDA_VISIBLE_DEVICES=$CUDA_DEVICES NPROC_PER_NODE=$NPROC_PER_NODE ENABLE_AUDIO_OUTPUT=False megatron sft \
    --model_type qwen3_omni \
    --load "$BASE_MODEL-mcore" \
    --no_load_optim false \
    --no_load_rng false \
    --moe_permute_fusion true \
    --attn_impl sdpa \
    --shuffle_buffer_size 100000 \
    --moe_grouped_gemm true \
    --moe_shared_expert_overlap true \
    --moe_aux_loss_coeff 1e-3 \
    --tensor_model_parallel_size 1 \
    --sequence_parallel false \
    --expert_model_parallel_size 2 \
    --packing true \
    --streaming true \
    --freeze_llm false \
    --freeze_vit true \
    --freeze_aligner true \
    --split_dataset_ratio 0.0 \
    --micro_batch_size 16 \
    --global_batch_size 64 \
    --recompute_granularity full \
    --recompute_method uniform \
    --recompute_num_layers 1 \
    --finetune true \
    --cross_entropy_loss_fusion true \
    --lr 1e-5 \
    --lr_warmup_fraction 0.12 \
    --min_lr 1e-6 \
    --max_epochs $STAGE1_EPOCHS \
    --train_iters 20000 \
    --eval_iters 40 \
    --save_interval 200 \
    --save $OUTPUT_PATH \
    --log_interval 10 \
    --tensorboard_dir ./tensorborad_logs/ \
    --vit_gradient_checkpointing true \
    --num_workers 8 \
    --no_save_optim false \
    --no_save_rng false \
    --dataset_num_proc 8 \
    --dataset $DATASET_PATH \
    --val_dataset $VAL_DATASET \
    --train_type  lora \
    --lora_rank 8 \
    --lora_alpha 32 \
    --target_modules all-linear \
    --max_length $STAGE1_MAX_LEN \
    --wandb_exp_name $WANDB_NAME \
    --wandb_project $WANDB_PROJECT \
    --wandb_save_dir wandb \
    --attention_backend flash \
    --system $DEFAULT_SYSTEM 

Checkpoint Structure

My checkpoint directory contains:

v20-20250927-180730/
├── iter_0000200/
│   ├── [checkpoint files]
├── latest_checkpointed_iteration.txt
└── [other metadata files]

Attempted Solutions

1. Loading from specific iteration directory

--adapter_load /path/to/checkpoint/v20-20250927-180730/iter_0000200

Result: KeyError: "language_model.decoder.layers.self_attention.linear_proj.weight from model not in state dict"
The checkpoint contains only LoRA weights (.lora_A.default.weight, .lora_B.default.weight) but the system expects full model weights.

2. Loading from checkpoint directory directly

--load /path/to/checkpoint/v20-20250927-180730

Result: Same KeyError - attempts to load LoRA checkpoint as full model.

3. Base model + adapter approach (current)

--load /path/to/base-model-mcore \
--adapter_load /path/to/checkpoint/v20-20250927-180730

Result: Training starts but from iteration 1, not 201. Logs show:

WARNING: could not find the metadata file .../latest_checkpointed_iteration.txt
Setting rerun_state_machine.current_iteration to 0...

4. Tried with --ckpt_step parameter

--ckpt_step 200

Result: ValueError: remaining_argv: ['--ckpt_step', '200'] - parameter not recognized.

Expected Behavior

• Resume training from iteration 201
• Restore optimizer state (Adam momentum, learning rate scheduler state)
• Continue with trained LoRA adapter weights

Current Behavior

• Training starts from iteration 1
• Learning rate starts from warmup phase (2.500000E-09) instead of where it should be after 200 iterations
• LoRA weights appear to be loaded correctly
• Optimizer state appears to be reset

Questions

1. What is the correct way to resume LoRA training in ms-swift with Megatron?
2. Are there specific parameters for resuming LoRA checkpoints vs full model checkpoints?
3. Should the checkpoint structure be different for proper resuming?
4. Is there a way to specify the starting iteration when resuming?

Additional Context

The initial training runs successfully and saves checkpoints every 200 iterations. The issue only occurs when trying to resume from an existing checkpoint. The LoRA adapter weights seem to be loaded correctly (model capabilities are maintained), but the training state (iteration count, optimizer, scheduler) is not properly restored.

Any guidance on the proper procedure for LoRA checkpoint resuming would be greatly appreciated.

Start by reproducing the resume command with --load pointing to the base model and --adapter_load pointing to iter_0000200, then inspect how latest_checkpointed_iteration.txt is discovered and how the Megatron checkpoint state is loaded. Done means the run resumes at iteration 201 with the LoRA weights, optimizer state, and scheduler state restored.