Error: Qwen3.5-35B-A3B lora sft with mcore-adapter, error occurs when saving a ckpt.

**training scripts:**
`# workdir=$(cd $(dirname $0); pwd)
# parent_dir=$(dirname "$workdir")

WORLD_SIZE=4
TENSOR_MODEL_PARALLEL_SIZE=1
EXPERT_MODEL_PARALLEL_SIZE=4
PIPELINE_MODEL_PARALLEL_SIZE=1

MODEL_NAME="/workspace/mount/Models/Qwen/Qwen3.5-35B-A3B"

export DISABLE_VERSION_CHECK=1

USE_MCA=true

mca_options=" \
--tensor_model_parallel_size ${TENSOR_MODEL_PARALLEL_SIZE} \
--sequence_parallel \
--pipeline_model_parallel_size ${PIPELINE_MODEL_PARALLEL_SIZE} \
--expert_model_parallel_size ${EXPERT_MODEL_PARALLEL_SIZE} \
--use_distributed_optimizer \
--bias_activation_fusion \
--apply_rope_fusion \
--overlap_param_gather \
--overlap_grad_reduce \
--moe_grouped_gemm true\
--moe_token_dispatcher_type alltoall"

llama_factory_options=" \
--deepspeed=./config/ds_zero2.json"

options=" \
--do_train \
--stage=sft \
--finetuning_type=lora \
--lora_rank 64 \
--dataset_dir=./mcore_adapter/examples/data \
--dataset=tulu-if-3w-alc \
--preprocessing_num_workers=8 \
--cutoff_len=2048 \
--template=qwen3_5 \
--model_name_or_path=$MODEL_NAME \
--output_dir=/workspace/mount/output3 \
--per_device_train_batch_size=1 \
--gradient_accumulation_steps=2 \
--calculate_per_token_loss=True \
--num_train_epochs 3 \
--learning_rate=2e-5 \
--logging_steps=1 \
--save_steps=50 \
--save_strategy steps \
--lr_scheduler_type=cosine \
--bf16 \
--recompute_granularity=full \
--max_steps=2"

if [ "$USE_MCA" = true ]; then
options="$options $mca_options --use_mca"
else
WORLD_SIZE=$(($WORLD_SIZE / $TENSOR_MODEL_PARALLEL_SIZE / $PIPELINE_MODEL_PARALLEL_SIZE))
options="$options $llama_factory_options --use_mca=False"
fi

torchrun --nproc_per_node=$WORLD_SIZE ./mcore_adapter/examples/train/run_train.py $options
`

---------------------------------------------------------------------------------------------
**Error Logs**

[rank3]: Traceback (most recent call last):
[rank3]: File "/mnt/home/ROLL/./mcore_adapter/examples/train/run_train.py", line 325, in
[rank3]: main()
[rank3]: File "/mnt/home/ROLL/./mcore_adapter/examples/train/run_train.py", line 319, in main
[rank3]: mca_train(training_args, model_args, data_args, finetuning_args)
[rank3]: File "/mnt/home/ROLL/./mcore_adapter/examples/train/run_train.py", line 292, in mca_train
[rank3]: sft_mca_train(training_args, model_args, data_args, finetuning_args)
[rank3]: File "/mnt/home/ROLL/./mcore_adapter/examples/train/run_train.py", line 232, in sft_mca_train
[rank3]: trainer.train(training_args.resume_from_checkpoint)
[rank3]: File "/usr/local/lib/python3.12/site-packages/transformers/trainer.py", line 1412, in train
[rank3]: return inner_training_loop(
[rank3]: ^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/mnt/home/ROLL/mcore_adapter/src/mcore_adapter/trainer/trainer.py", line 889, in _inner_training_loop
[rank3]: self._maybe_log_save_evaluate(
[rank3]: File "/mnt/home/ROLL/mcore_adapter/src/mcore_adapter/trainer/trainer.py", line 1015, in _maybe_log_save_evaluate
[rank3]: self._save_checkpoint(model, trial)
[rank3]: File "/usr/local/lib/python3.12/site-packages/transformers/trainer.py", line 3058, in _save_checkpoint
[rank3]: self._save_optimizer_and_scheduler(output_dir)
[rank3]: File "/mnt/home/ROLL/mcore_adapter/src/mcore_adapter/trainer/trainer.py", line 1067, in _save_optimizer_and_scheduler
[rank3]: model_shared_state_dict = self.model.sharded_state_dict()
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/mnt/home/ROLL/mcore_adapter/src/mcore_adapter/models/model_factory.py", line 193, in sharded_state_dict
[rank3]: state_dict["model"] = self.models[0].sharded_state_dict(prefix, *args, **kwargs)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/models/gpt/gpt_model.py", line 819, in sharded_state_dict
[rank3]: sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/models/common/language_module/language_module.py", line 290, in sharded_state_dict
[rank3]: sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/module.py", line 99, in sharded_state_dict
[rank3]: sharded_state_dict_default(
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/utils.py", line 241, in sharded_state_dict_default
[rank3]: module_sharded_sd = module.sharded_state_dict(
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/transformer_block.py", line 874, in sharded_state_dict
[rank3]: layer_sharded_state_dict = layer.sharded_state_dict(
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/transformer_layer.py", line 851, in sharded_state_dict
[rank3]: sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/module.py", line 99, in sharded_state_dict
[rank3]: sharded_state_dict_default(
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/transformer/utils.py", line 241, in sharded_state_dict_default
[rank3]: module_sharded_sd = module.sharded_state_dict(
[rank3]: ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: File "/usr/local/lib/python3.12/site-packages/megatron/core/ssm/gated_delta_net.py", line 470, in sharded_state_dict
[rank3]: assert sharded_state_dict[f"{prefix}in_proj.weight"].data.size(0) == in_proj_dim_local_tp, (
[rank3]: ~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: KeyError: 'decoder.layers.0.self_attention.in_proj.weight'

Reproduce the failure with the provided torchrun script and save_steps configuration. Start in src/mcore_adapter/trainer/trainer.py at _save_optimizer_and_scheduler, then trace model_factory.py and the Megatron gated_delta_net.py sharded_state_dict call. Done means the Qwen3.5-35B-A3B LoRA run saves a checkpoint without the missing in_proj.weight KeyError.

Written by the indexing model from the issue text.