All sglang TP workers process hidden states unnecessarily, causing training pipeline slowdown

[CH3CHOHCH3]

Background

I noticed that the training speed was very slow when I was training a Dflash draft model for Minimax-m2.5 on a 8xH200 server, taking about 100–200 seconds per step.

I allocated 4 GPUs to sglang TP workers and 4 GPUs to training.

My dataset's context length was around 32k.

Root Cause

In fact, the GPUs' utilization was 0% most of the time, which suggests that something may be stuck on the CPU.

So I checked the output of top and found 3 sglang tp workers has 100% CPU core utilizatios.

Then I used py-spy dump to inspect their python stack and found that they spent most time on _process_hidden_states_for_req:

TorchSpec/patches/sglang/v0.5.12/sglang.patch

Lines 590 to 618 in d230a3f

	+ def _process_hidden_states_for_req(
	+ self: Scheduler,
	+ req: Req,
	+ batch: ScheduleBatch,
	+ logits_output: LogitsProcessorOutput,
	+ hidden_state_offset: int,
	+ copy_done_event=None,
	+ ):
	+ """Process hidden states during prefill for spec training or return_hidden_states."""
	+ seq_len = len(req.origin_input_ids)
	+ req_hidden_states = logits_output.hidden_states[
	+ hidden_state_offset : hidden_state_offset + seq_len
	+ ]
	+
	+ if (
	+ batch.spec_training_info is not None
	+ and batch.spec_training_info.has_request(req.rid)
	+ and self.eagle_mooncake_store is not None
	+ ):
	+ self._send_hidden_states_to_mooncake(
	+ req, batch, req_hidden_states, logits_output, hidden_state_offset,
	+ copy_done_event=copy_done_event,
	+ )
	+ else:
	+ if copy_done_event is not None:
	+ copy_done_event.synchronize()
	+ req.hidden_states.append(
	+ req_hidden_states.cpu().clone().tolist()
	+ )

I thought they were doingreq_hidden_states.cpu().clone().tolist() because the eagle_mooncake_store is not None only for TP rank 0.

If TP worker 0 can put all hidden states to mooncake, the other TP workers' should_process_hidden_states should be false:

TorchSpec/patches/sglang/v0.5.12/sglang.patch

Lines 567 to 574 in d230a3f

	+ should_process_hidden_states = (
	+ logits_output.hidden_states is not None
	+ and (
	+ req.return_hidden_states
	+ or (
	+ req.spec_training_data_id is not None
	+ and self.attn_tp_rank == 0
	+ )

But sgl_engine.py set return_hidden_states to Ture so all TP workers will enter _process_hidden_states_for_req which is likely not expected:

TorchSpec/torchspec/inference/engine/sgl_engine.py

Lines 428 to 433 in d230a3f

	engine_kwargs = {
	"prompt": formatted_prompts,
	"spec_training_data_id": data_ids,
	"sampling_params": {"max_new_tokens": 0},
	"return_hidden_states": True,
	}

Fix and Test

I removed "return_hidden_states": True in sgl_engine.py and restarted the train process, the speed was improved to 10s/step.

Summary

The current hidden states processing logic appears unreasonable:

1. "return_hidden_states": True will let all TP workers enter _process_hidden_states_for_req;
2. But only TP rank 0 has eagle_mooncake_store.
3. So other TP workers will cost lots of time on req_hidden_states.cpu().clone().tolist(), which is meaningless.