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

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

Start with scripts/train_dflash.py and trace the failing forward pass through specforge/core/dflash.py into specforge/modeling/draft/dflash.py, especially the fc call reported in the traceback. Run the qwen3.5-35b-a3b reproduction and inspect the hidden-state and linear-layer shapes. Done means training no longer fails with the matrix-shape error and results are obtained beyond one target-model layer.