Megatron GKD OOM when compute_teacher_logits

Describe the bug
16 * A800 80GB, megatron gkd, teacher_model: 235B-A22B，model: 30B-A3B
OOM when compute_teacher_logits

Additional context

megatron rlhf \
    --rlhf_type gkd \
    --model ${MODEL} \
    --teacher_model ${TEACHER_MODEL} \
    --dataset ${DATASETS} \
    --train_type full \
    --load_safetensors true \
    --save_safetensors true \
    --split_dataset_ratio 0 \
    --pipeline_model_parallel_size 2 \
    --expert_model_parallel_size 8 \
    --moe_permute_fusion true \
    --moe_grouped_gemm true \
    --moe_shared_expert_overlap true \
    --moe_aux_loss_coeff 1e-3 \
    --micro_batch_size 1 \
    --global_batch_size 32 \
    --offload_teacher_model true \
    --packing false \
    --recompute_granularity full \
    --recompute_method uniform \
    --recompute_num_layers 1 \
    --max_epochs 3 \
    --finetune \
    --cross_entropy_loss_fusion true \
    --lr 5e-5 \
    --use_vllm true \
    --vllm_mode colocate \
    --vllm_max_model_len 16000 \
    --vllm_gpu_memory_utilization 0.15 \
    --vllm_tensor_parallel_size 8 \
    --vllm_enforce_eager true \
    --max_completion_length 8192 \
    --lr_warmup_fraction 0.05 \
    --save ${OUTPUT_MODEL} \
    --loss_scale last_round \
    --num_workers 8 \
    --dataset_num_proc 8 \
    --no_save_optim true \
    --no_save_rng true \
    --temperature 1.0 \
    --sequence_parallel true \
    --attention_backend flash \
    --log_interval 1 

2026-01-06 11:41:57
Variable._execution_engine.run_backward(
2026-01-06 11:42:15
[2026-01-06 11:42:15] iteration 1/ 382 | consumed samples: 32 | elapsed time per iteration (ms): 198314.0 | memory(GiB): 67.84 | elapsed time: 3m 18s | remaining time: 20h 59m 17s | learning rate: 2.617801E-06 | global batch size: 32 | jsd_loss: 7.980616E-02 | loss: 7.980616E-02 | load_balancing_loss: 1.637991E+00 | loss scale: 1.0 | grad norm: 1.165 | number of skipped iterations: 0 | number of nan iterations: 0 |
2026-01-06 11:42:15
[Rank 11] (after 1 iterations) memory (MB) | allocated: 45670.0634765625 | max allocated: 69631.9677734375 | reserved: 45782.0 | max reserved: 70022.0
2026-01-06 11:42:15
[Rank 9] (after 1 iterations) memory (MB) | allocated: 45670.0634765625 | max allocated: 69631.9677734375 | reserved: 45762.0 | max reserved: 70022.0
2026-01-06 11:42:15
[Rank 10] (after 1 iterations) memory (MB) | allocated: 45670.0634765625 | max allocated: 69631.9677734375 | reserved: 45762.0 | max reserved: 70002.0
2026-01-06 11:42:15
[Rank 8] (after 1 iterations) memory (MB) | allocated: 45670.0634765625 | max allocated: 69631.9677734375 | reserved: 45782.0 | max reserved: 69982.0
2026-01-06 11:42:34
INFO 01-06 11:42:34 [block_pool.py:378] Successfully reset prefix cache
2026-01-06 11:42:34
INFO 01-06 11:42:34 [block_pool.py:378] Successfully reset prefix cache
2026-01-06 11:42:34
INFO 01-06 11:42:34 [block_pool.py:378] Successfully reset prefix cache
2026-01-06 11:42:34
INFO 01-06 11:42:34 [block_pool.py:378] Successfully reset prefix cache
2026-01-06 11:42:34
INFO 01-06 11:42:34 [block_pool.py:378] Successfully reset prefix cache
2026-01-06 11:42:34
INFO 01-06 11:42:34 [block_pool.py:378] Successfully reset prefix cache
2026-01-06 11:42:34
INFO 01-06 11:42:34 [block_pool.py:378] Successfully reset prefix cache
2026-01-06 11:42:34
INFO 01-06 11:42:34 [block_pool.py:378] Successfully reset prefix cache
2026-01-06 11:42:34
[WARNING:swift] max_model_len(16000) - num_tokens(9067) < max_tokens(8192). Setting max_tokens: 6933
2026-01-06 11:43:21
[rank9]: Traceback (most recent call last):
2026-01-06 11:43:21
[rank9]: File "/workspace/ms-swift/swift/cli/_megatron/rlhf.py", line 7, in <module>
2026-01-06 11:43:21
[rank9]: megatron_rlhf_main()
2026-01-06 11:43:21
[rank9]: File "/workspace/ms-swift/swift/megatron/train/rlhf.py", line 78, in megatron_rlhf_main
2026-01-06 11:43:21
[rank9]: return MegatronRLHF(args).main()
2026-01-06 11:43:21
[rank9]: File "/workspace/ms-swift/swift/llm/base.py", line 49, in main
2026-01-06 11:43:21
[rank9]: result = self.run()
2026-01-06 11:43:21
[rank9]: File "/workspace/ms-swift/swift/megatron/train/sft.py", line 77, in run
2026-01-06 11:43:21
[rank9]: self.trainer.train(train_dataset, val_dataset, data_collator)
2026-01-06 11:43:21
[rank9]: File "/workspace/ms-swift/swift/megatron/trainers/gkd_trainer.py", line 80, in train
2026-01-06 11:43:21
[rank9]: super().train(train_dataset, val_dataset, data_collator)
2026-01-06 11:43:21
[rank9]: File "/workspace/ms-swift/swift/megatron/trainers/base.py", line 1098, in train
2026-01-06 11:43:21
[rank9]: pretrain(
2026-01-06 11:43:21
[rank9]: File "/workspace/Megatron-LM/megatron/training/training.py", line 737, in pretrain
2026-01-06 11:43:21
[rank9]: iteration, num_floating_point_operations_so_far = train(
2026-01-06 11:43:21
[rank9]: File "/workspace/Megatron-LM/megatron/training/training.py", line 2298, in train
2026-01-06 11:43:21
[rank9]: ) = train_step(
2026-01-06 11:43:21
[rank9]: File "/workspace/ms-swift/swift/megatron/trainers/base.py", line 564, in train_step
2026-01-06 11:43:21
[rank9]: new_data_iterator = self._replace_data_iterator(data_iterator, model)
2026-01-06 11:43:21
[rank9]: File "/workspace/ms-swift/swift/megatron/trainers/gkd_trainer.py", line 484, in _replace_data_iterator
2026-01-06 11:43:21
[rank9]: self._compute_teacher_logits(encoded_batches)
2026-01-06 11:43:21
[rank9]: File "/workspace/ms-swift/swift/megatron/trainers/gkd_trainer.py", line 441, in _compute_teacher_logits
2026-01-06 11:43:21
[rank9]: teacher_logits = forward_step_helper(teacher_model, teacher_data)
2026-01-06 11:43:21
[rank9]: File "/workspace/ms-swift/swift/megatron/utils/utils.py", line 310, in forward_step_helper
2026-01-06 11:43:21
[rank9]: output_tensor = model(**inputs)
2026-01-06 11:43:21
[rank9]: File "/root/anaconda3/envs/swift/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
2026-01-06 11:43:21
[rank9]: return self._call_impl(*args, **kwargs)
2026-01-06 11:43:21
[rank9]: File "/root/anaconda3/envs/swift/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
2026-01-06 11:43:21
[rank9]: return forward_call(*args, **kwargs)
2026-01-06 11:43:21
[rank9]: File "/workspace/Megatron-LM/megatron/core/transformer/module.py", line 434, in forward
2026-01-06 11:43:21
[rank9]: outputs = float16_to_fp32(outputs)
2026-01-06 11:43:21
[rank9]: File "/workspace/Megatron-LM/megatron/core/transformer/module.py", line 362, in float16_to_fp32
2026-01-06 11:43:21
[rank9]: return conversion_helper(val, float_conversion)
2026-01-06 11:43:21
[rank9]: File "/workspace/Megatron-LM/megatron/core/transformer/module.py", line 321, in conversion_helper
2026-01-06 11:43:21
[rank9]: return conversion(val)
2026-01-06 11:43:21
[rank9]: File "/workspace/Megatron-LM/megatron/core/transformer/module.py", line 359, in float_conversion
2026-01-06 11:43:21
[rank9]: val = val.float()
2026-01-06 11:43:21
[rank9]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 586.00 MiB. GPU 1 has a total capacity of 79.35 GiB of which 420.12 MiB is free. Process 2903556 has 78.86 GiB memory in use. Of the allocated memory 73.72 GiB is allocated by PyTorch, and 129.11 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

Start in swift/megatron/trainers/gkd_trainer.py at _compute_teacher_logits and trace the call through swift/megatron/utils/utils.py. Compare that path with Megatron-LM's megatron/core/transformer/module.py, especially float16_to_fp32, and reproduce the provided megatron rlhf command. Done means teacher-logit computation completes without the reported CUDA out-of-memory error under the shown configuration.