[Bug] qwen3.5-35b-a3b dflash

[Jim2016713]

Checklist

• 1. I have searched related issues but cannot get the expected help.
• 2. The bug has not been fixed in the latest version.
• 3. Please note that if the bug-related issue you submitted lacks corresponding environment info and a minimal reproducible demo, it will be challenging for us to reproduce and resolve the issue, reducing the likelihood of receiving feedback.
• 4. If the issue you raised is not a bug but a question, please raise a discussion at https://github.com/sgl-project/SpecForge/discussions/new/choose Otherwise, it will be closed.
• 5. Please use English, otherwise it will be closed.

Describe the bug

only obtained the results from one layer of the target model

[rank5]: RuntimeError: mat1 and mat2 shapes cannot be multiplied (5900x2048 and 10240x2048)
[rank2]: Traceback (most recent call last):
[rank2]: File "/dc-hl/zibing.wei/code/dflash/qwen3.5/SpecForge/scripts/train_dflash.py", line 592, in
[rank2]: main()
[rank2]: File "/dc-hl/zibing.wei/code/dflash/qwen3.5/SpecForge/scripts/train_dflash.py", line 534, in main
[rank2]: loss, accuracy = dflash_model(
[rank2]: ^^^^^^^^^^^^^
[rank2]: File "/workspace/SpecForge/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
[rank2]: return self._call_impl(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/workspace/SpecForge/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
[rank2]: return forward_call(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/workspace/SpecForge/.venv/lib/python3.11/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 851, in forward
[rank2]: output = self._fsdp_wrapped_module(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/workspace/SpecForge/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
[rank2]: return self._call_impl(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/workspace/SpecForge/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
[rank2]: return forward_call(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/dc-hl/zibing.wei/code/dflash/qwen3.5/SpecForge/specforge/core/dflash.py", line 218, in forward
[rank2]: output_hidden = self.draft_model(
[rank2]: ^^^^^^^^^^^^^^^^^
[rank2]: File "/workspace/SpecForge/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
[rank2]: return self._call_impl(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/workspace/SpecForge/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
[rank2]: return forward_call(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/dc-hl/zibing.wei/code/dflash/qwen3.5/SpecForge/specforge/modeling/draft/dflash.py", line 253, in forward
[rank2]: target_hidden = self.hidden_norm(self.fc(target_hidden))
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/workspace/SpecForge/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
[rank2]: return self._call_impl(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/workspace/SpecForge/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
[rank2]: return forward_call(*args, **kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/workspace/SpecForge/.venv/lib/python3.11/site-packages/torch/nn/modules/linear.py", line 134, in forward
[rank2]: return F.linear(input, self.weight, self.bias)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: RuntimeError: mat1 and mat2 shapes cannot be multiplied (6219x2048 and 10240x2048)

Reproduction

qwen3.5-35b-a3b

Environment

qwen3.5-35b-a3b

[Mandy3311]

Hi, you'll need to add DFlash support in SGLang first. You could reference this commit: sgl-project/sglang@1a240a2. should work as expected once the inference backend DFlash support integrated.

[Jim2016713]

Hi, you'll need to add DFlash support in SGLang first. You could reference this commit: sgl-project/sglang@1a240a2. should work as expected once the inference backend DFlash support integrated.

I reinstalled it using your repository branch and got this error.
[rank2]: Traceback (most recent call last):
[rank2]: File "/dc-hl/zibing.wei/code/dflash/qwen3.5/SpecForge/scripts/train_dflash.py", line 591, in
[rank2]: main()
[rank2]: File "/dc-hl/zibing.wei/code/dflash/qwen3.5/SpecForge/scripts/train_dflash.py", line 402, in main
[rank2]: target_model, draft_model = build_models(args)
[rank2]: ^^^^^^^^^^^^^^^^^^
[rank2]: File "/dc-hl/zibing.wei/code/dflash/qwen3.5/SpecForge/scripts/train_dflash.py", line 160, in build_models
[rank2]: target_model = get_dflash_target_model(
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/dc-hl/zibing.wei/code/dflash/qwen3.5/SpecForge/specforge/modeling/target/dflash_target_model.py", line 299, in get_dflash_target_model
[rank2]: return SGLangDFlashTargetModel.from_pretrained(
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/dc-hl/zibing.wei/code/dflash/qwen3.5/SpecForge/specforge/modeling/target/dflash_target_model.py", line 83, in from_pretrained
[rank2]: server_args = ServerArgs(
[rank2]: ^^^^^^^^^^^
[rank2]: TypeError: ServerArgs.init() got an unexpected keyword argument 'enable_piecewise_cuda_graph'

[Jim2016713]

Hi, you'll need to add DFlash support in SGLang first. You could reference this commit: sgl-project/sglang@1a240a2. should work as expected once the inference backend DFlash support integrated.

Hello, have you encountered this issue where after the model and data have been loaded, the training process gets stuck for a long time with zero GPU utilization?