modelscope / modelscope/ms-swift

validation invalid input shape error while using cp>1 for megatron swift

Open
#6,195 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement stale
Dominant language
Python
Stars
15.7k
Forks
1.7k
Avg merge
1d 16h
Merged PRs (30d)
136

Description

Describe the bug
My training script is given as following:

PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
megatron sft \
    --finetune true \
    --train_type lora \
    --lora_rank 8 \
    --lora_alpha 32 \
    --target_modules all-linear \
    --load $LOAD_DIR \
    --cached_dataset $MODELSCOPE_CACHE/tokenized-$MODEL_TYPE-$STRIPPED_DATASET \
    --split_dataset_ratio $SPLIT_DATASET_RATIO \
    --dataset_num_proc $DATASET_NUM_PROC \
    --packing true \
    --context-parallel-size 4 \
    --tensor_model_parallel_size 8 \
    --pipeline_model_parallel_size 1 \
    --micro_batch_size 1 \
    --global_batch_size 128 \
    --max_length $MAX_LEN \
    --recompute_granularity selective \
    --overlap_param_gather true \
    --overlap_grad_reduce true \
    --cross_entropy_loss_fusion true \
    --lr 5e-6 \
    --lr_decay_style cosine \
    --lr_warmup_fraction 0.05 \
    --min_lr 0 \
    --max_epochs 5 \
    --save $SAVE_DIR \
    --tensorboard_dir $SAVE_DIR/tf_logs \
    --tensorboard_log_interval 1 \
    --log_interval 1 \
    --save_interval 100 \
    --eval_interval 50 \
    --num_workers $DATASET_NUM_PROC \
    --attention_backend flash

The full error log:

2025-10-18 01:36:06.294112153  nm129-a100-80g-46 ['Traceback (most recent call last):\n', '  File "/usr/local/lib/python3.10/dist-packages/swift/megatron/trainers/trainer.py", line 100, in forward_step\n    data = get_batch(data_iterator)\n', '  File "/usr/local/lib/python3.10/dist-packages/swift/megatron/trainers/utils.py", line 205, in get_batch\n    batch = get_batch_on_this_cp_rank(batch)\n', '  File "/usr/local/lib/python3.10/dist-packages/swift/megatron/trainers/utils.py", line 192, in get_batch_on_this_cp_rank\n    batch[key] = _split_tokens(val, packed_seq_params.cu_seqlens_q)\n', '  File "/usr/local/lib/python3.10/dist-packages/swift/megatron/trainers/utils.py", line 160, in _split_tokens\n    val = val.view(\n', "RuntimeError: shape '[1, 8, 4095]' is invalid for input of size 32766\n"]
2025-10-18 01:36:06.294513581  nm129-a100-80g-46 [INFO:swift] images_dir: /mnt/ht1-g1-gpfs/ckpts/Qwen3-32B/cpt/custom-cpt-0924/sft_lora/liucong-sampled-data_ht-data-hmm-1011-packing/v0-20251017-173520/images
2025-10-18 01:36:06.509606752  nm129-a100-80g-46 INFO:matplotlib.font_manager:generated new fontManager
2025-10-18 01:36:06.777615428  nm129-a100-80g-46 Traceback (most recent call last):
2025-10-18 01:36:06.777646146  nm129-a100-80g-46   File "/usr/local/lib/python3.10/dist-packages/swift/cli/_megatron/sft.py", line 5, in <module>
2025-10-18 01:36:06.777653446  nm129-a100-80g-46     megatron_sft_main()
2025-10-18 01:36:06.777660716  nm129-a100-80g-46   File "/usr/local/lib/python3.10/dist-packages/swift/megatron/train/sft.py", line 72, in megatron_sft_main
2025-10-18 01:36:06.777666845  nm129-a100-80g-46     return MegatronSft(args).main()
025-10-18 01:36:06.777672867  nm129-a100-80g-46   File "/usr/local/lib/python3.10/dist-packages/swift/llm/base.py", line 49, in main
2025-10-18 01:36:06.777679375  nm129-a100-80g-46     result = self.run()
2025-10-18 01:36:06.777685186  nm129-a100-80g-46   File "/usr/local/lib/python3.10/dist-packages/swift/megatron/train/sft.py", line 62, in run
2025-10-18 01:36:06.777691319  nm129-a100-80g-46     self.trainer.train(train_dataset, val_dataset, data_collator)
2025-10-18 01:36:06.777697092  nm129-a100-80g-46   File "/usr/local/lib/python3.10/dist-packages/swift/megatron/trainers/base.py", line 415, in train
2025-10-18 01:36:06.777703181  nm129-a100-80g-46     pretrain(
2025-10-18 01:36:06.777709499  nm129-a100-80g-46   File "/mnt/ht1-g1-gpfs/lujzh/ms-swift/Megatron-LM/megatron/training/training.py", line 864, in pretrain
2025-10-18 01:36:06.777715710  nm129-a100-80g-46     iteration, num_floating_point_operations_so_far = train(
2025-10-18 01:36:06.777721817  nm129-a100-80g-46   File "/mnt/ht1-g1-gpfs/lujzh/ms-swift/Megatron-LM/megatron/training/training.py", line 2279, in train
2025-10-18 01:36:06.777727585  nm129-a100-80g-46     ) = train_step(
2025-10-18 01:36:06.777733344  nm129-a100-80g-46   File "/usr/local/lib/python3.10/dist-packages/swift/megatron/trainers/base.py", line 245, in train_step
2025-10-18 01:36:06.777739178  nm129-a100-80g-46     return self._origin_train_step(forward_step_func, new_data_iterator, model, optimizer, opt_param_scheduler,
2025-10-18 01:36:06.777745006  nm129-a100-80g-46   File "/mnt/ht1-g1-gpfs/lujzh/ms-swift/Megatron-LM/megatron/training/training.py", line 1395, in train_step
2025-10-18 01:36:06.777750736  nm129-a100-80g-46     losses_reduced = forward_backward_func(
2025-10-18 01:36:06.777757196  nm129-a100-80g-46   File "/mnt/ht1-g1-gpfs/lujzh/ms-swift/Megatron-LM/megatron/core/pipeline_parallel/schedules.py", line 500, in forward_backward_no_pipelining
2025-10-18 01:36:06.777763198  nm129-a100-80g-46     output_tensor, num_tokens = forward_step(
2025-10-18 01:36:06.777788069  nm129-a100-80g-46   File "/mnt/ht1-g1-gpfs/lujzh/ms-swift/Megatron-LM/megatron/core/pipeline_parallel/schedules.py", line 289, in forward_step
2025-10-18 01:36:06.777794048  nm129-a100-80g-46     output_tensor, loss_func = forward_step_func(data_iterator, model)
2025-10-18 01:36:06.777799797  nm129-a100-80g-46   File "/usr/local/lib/python3.10/dist-packages/swift/megatron/trainers/trainer.py", line 100, in forward_step
2025-10-18 01:36:06.777805555  nm129-a100-80g-46     data = get_batch(data_iterator)
2025-10-18 01:36:06.777811332  nm129-a100-80g-46   File "/usr/local/lib/python3.10/dist-packages/swift/megatron/trainers/utils.py", line 205, in get_batch
2025-10-18 01:36:06.777817050  nm129-a100-80g-46     batch = get_batch_on_this_cp_rank(batch)
2025-10-18 01:36:06.777823056  nm129-a100-80g-46   File "/usr/local/lib/python3.10/dist-packages/swift/megatron/trainers/utils.py", line 192, in get_batch_on_this_cp_rank
2025-10-18 01:36:06.777828750  nm129-a100-80g-46     batch[key] = _split_tokens(val, packed_seq_params.cu_seqlens_q)
2025-10-18 01:36:06.777834599  nm129-a100-80g-46   File "/usr/local/lib/python3.10/dist-packages/swift/megatron/trainers/utils.py", line 160, in _split_tokens
2025-10-18 01:36:06.777840388  nm129-a100-80g-46     val = val.view(
2025-10-18 01:36:06.777846567  nm129-a100-80g-46 RuntimeError: shape '[1, 8, 4095]' is invalid for input of size 32766

Your hardware and system info
Write your system info like CUDA version/system/GPU/torch version here(在这里给出硬件信息和系统信息，如CUDA版本，系统，GPU型号和torch版本等)

I used 88A100 80g for training and the driver version 545.23.08. I list the other dependencies as following:

root@ji-jupyter-147108455869604416-master-0:/workspace# pip list | grep torch
pytorch-quantization          2.1.2
torch                         2.3.0a0+ebedce2
torch-tensorrt                2.3.0a0
torchdata                     0.7.1a0
torchtext                     0.17.0a0
torchvision                   0.18.0a0
root@ji-jupyter-147108455869604416-master-0:/workspace# pip list | grep trans
s3transfer                    0.13.1
sentence-transformers         5.1.0
transformer-engine            2.3.0+5de3e14
transformers                  4.55.4
transformers-stream-generator 0.0.5
root@ji-jupyter-147108455869604416-master-0:/workspace# pip list | grep swift
ms_swift                      3.7.3
root@ji-jupyter-147108455869604416-master-0:/workspace# nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Wed_Nov_22_10:17:15_PST_2023
Cuda compilation tools, release 12.3, V12.3.107
Build cuda_12.3.r12.3/compiler.33567101_0

Additional context
Add any other context about the problem here(在这里补充其他信息)

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in swift/megatron/trainers/utils.py at get_batch_on_this_cp_rank and _split_tokens, using the reported invalid view and packed_seq_params.cu_seqlens_q as the entry point. Reproduce the provided megatron sft command with context parallelism, then determine why the token shape and input size disagree; done means the command no longer raises the shape error and training proceeds.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.