[Question] Any plans to support CP for GDN layers?

[Andrewclorsy]

Your Question

Training qwen3.5-27B in long context will result in OOM during backpropagation. Any plans to support CP for GDN layers?

What I've Tried

Training qwen3.5-27B in long context will result in OOM during backpropagation. Any plans to support CP for GDN layers?

Environment (if relevant)

• slime version:
• Python version:
• PyTorch version:
• CUDA/ROCm version:
• GPU type and count:
• OS:

Additional Context

No response

Pre-submission Checklist

• I have read the CONTRIBUTING.md and understand the collaboration scope.
• I have read the documentation and FAQ and my question is not answered there.
• I have searched for existing issues and my question has not been asked before.

[zhuzilin]

Yes, this is on our roadmap.

To better understand the requirement and prioritize the implementation, could you share a bit more detail about your training setup?

• How many GPUs are you using?
• What context length are you targeting?

This would help us evaluate the demand and design the support more appropriately.

[Andrewclorsy]

Yes, this is on our roadmap.

To better understand the requirement and prioritize the implementation, could you share a bit more detail about your training setup?

• How many GPUs are you using?
• What context length are you targeting?

This would help us evaluate the demand and design the support more appropriately.

Thanks for the quick response — looking forward to it ！

Here are the details of my training setup:

Hardware:64 × 80GB GPUs (8 nodes × 8 GPUs/node)
Target context length: 128K

For Qwen3.5-27B ：
ARGS=(
   --tensor-model-parallel-size 4
   --pipeline-model-parallel-size 2
   --context-parallel-size 8
   --sequence-parallel
   --recompute-granularity full
   --recompute-method uniform
   --recompute-num-layers 1

   --rollout-batch-size 128
   --n-samples-per-prompt 8
   --rollout-max-response-len 128000
   --rollout-temperature 1
   --global-batch-size 1024
   --balance-data
   --group-rm
   --over-sampling-batch-size 160  
   --dynamic-sampling-filter-path slime.rollout.filter_hub.dynamic_sampling_filters.xxxxx

    --custom-generate-function-path xxxxx.generate
    --custom-config-path  xxxxx.yaml
）

With this setup, the forward and rollout work fine at 128K context. The OOM occurs specifically during backpropagation.
I also experimented with Qwen3.5-35B-A3B (MoE) using the same setup above with --expert-model-parallel-size 8. For Qwen3.5-35B-A3B, the GDN layers did not become a bottleneck and training ran normally at 128K context.

[lawrence-harmonic]

May be of interest:
NVIDIA/Megatron-LM#3282
NVIDIA/Megatron-LM#2642

[lawrence-harmonic]

Another reference impl: PrimeIntellect-ai/prime-rl#2080

[Zhichenzzz]

Hi @Andrewclorsy @zhuzilin, Miles has implemented the real CP for GDN layers on Qwen3.5 series. Maybe you can check it radixark/miles#885 as reference~

[hideonevan]

Yes, this is on our roadmap.
To better understand the requirement and prioritize the implementation, could you share a bit more detail about your training setup?

• How many GPUs are you using?
• What context length are you targeting?

This would help us evaluate the demand and design the support more appropriately.

Thanks for the quick response — looking forward to it ！

Here are the details of my training setup:

Hardware:64 × 80GB GPUs (8 nodes × 8 GPUs/node)
Target context length: 128K

For Qwen3.5-27B ：
ARGS=(
   --tensor-model-parallel-size 4
   --pipeline-model-parallel-size 2
   --context-parallel-size 8
   --sequence-parallel
   --recompute-granularity full
   --recompute-method uniform
   --recompute-num-layers 1

   --rollout-batch-size 128
   --n-samples-per-prompt 8
   --rollout-max-response-len 128000
   --rollout-temperature 1
   --global-batch-size 1024
   --balance-data
   --group-rm
   --over-sampling-batch-size 160  
   --dynamic-sampling-filter-path slime.rollout.filter_hub.dynamic_sampling_filters.xxxxx

    --custom-generate-function-path xxxxx.generate
    --custom-config-path  xxxxx.yaml
）

With this setup, the forward and rollout work fine at 128K context. The OOM occurs specifically during backpropagation. I also experimented with Qwen3.5-35B-A3B (MoE) using the same setup above with --expert-model-parallel-size 8. For Qwen3.5-35B-A3B, the GDN layers did not become a bottleneck and training ran normally at 128K context.

May I ask which version of Megatron-LM you're using?