Is this page helpful?

core.transformer.moe.router_trace#

Router decision tracing for MoE models for both training and inference.

Captures per-layer top-K routing decisions to a JSONL file for offline analysis of routing patterns (e.g., expert load balance, overlap between (layer N-2, N)). Enable via --moe-routing-trace-path in both training and inference.

Output format: one JSONL file per rank, one record per (step, block, layer): {“step”: 0, “stage”: “pre_dispatch”, “block”: “decoder”, “layer”: 3, “rank”: 0, “num_tokens”: 128, “topk”: 22, “top_indices”: [[12, 45, …], …]} MTP records carry an extra “mtp_idx” field so they never collide with decoder layers that share a layer number.

Optional sidecar binary files written:

  • hidden_states_rank{rank}.bin — bfloat16 hidden-state tensors; each JSONL record gains hs_offset, hs_bytes, hs_shape fields.

  • logits_rank{rank}.bin — bfloat16 pre-topk routing logits; each JSONL record gains logit_offset, logit_bytes, logit_shape fields.

Use load_hidden_states_for_record / load_logits_for_record to read sidecar tensors.

Note: Python forward hooks do not fire during CUDA graph replay. Run with --cuda-graph-impl none.

Module Contents#

Classes#

RouterTracer

Captures router top-K decisions across all MoE layers per step.

Functions#

_parse_router_module_name

Parse a router module name into (block, mtp_idx, layer).

init_moe_router_tracer

Initialize the global router tracer. Call after torch.distributed is initialized and before register_hooks is called on the model.

get_moe_router_tracer

Return the active tracer, or None if tracing is disabled.

load_hidden_states_for_record

Load the hidden-state tensor for a single JSONL record.

load_logits_for_record

Load the pre-topk routing logits for a single JSONL record.

Data#

_MOE_ROUTER_TRACER

_MTP_STACK_LAYER_RE

_MTP_LAYER_RE

_DECODER_LAYER_RE

API#

core.transformer.moe.router_trace._MOE_ROUTER_TRACER: Optional[core.transformer.moe.router_trace.RouterTracer]#

None

core.transformer.moe.router_trace._MTP_STACK_LAYER_RE#

‘compile(…)’

core.transformer.moe.router_trace._MTP_LAYER_RE#

‘compile(…)’

core.transformer.moe.router_trace._DECODER_LAYER_RE#

‘compile(…)’

core.transformer.moe.router_trace._parse_router_module_name(
module_name: str,
) → Optional[Tuple[str, Optional[int], int]]#

Parse a router module name into (block, mtp_idx, layer).

Returns None if the name matches neither the decoder nor the MTP pattern.

.. rubric:: Examples

decoder.layers.3.mlp.router -> (“decoder”, None, 3) mtp.layers.0.mtp_model_layer.layers.1.mlp.router -> (“mtp”, 0, 1) mtp.layers.0.mtp_model_layer.mlp.router -> (“mtp”, 0, 0)

core.transformer.moe.router_trace.init_moe_router_tracer(
output_dir: str,
max_steps: int,
rank: int,
training_mode: bool = False,
capture_hidden_states: bool = False,
capture_logits: bool = False,
dump_router_weights: bool = False,
) → None#

Initialize the global router tracer. Call after torch.distributed is initialized and before register_hooks is called on the model.

Parameters:

  • output_dir – Directory for JSONL trace files (and optional sidecars).

  • max_steps – Maximum steps (iterations in training, decode steps in inference) to capture.

  • rank – Distributed rank.

  • training_mode – If True, step boundaries are driven by advance_step() calls from the training

  • inference. (loop rather than the layer-repeat heuristic used during)

  • capture_hidden_states – Capture the input hidden-state tensor for each router call.

  • capture_logits – Capture pre-topk routing logits.

  • dump_router_weights – Save router weight tensors to a .pt file.

core.transformer.moe.router_trace.get_moe_router_tracer() → Optional[RouterTracer]#

Return the active tracer, or None if tracing is disabled.

core.transformer.moe.router_trace.load_hidden_states_for_record(
record: dict,
trace_dir: str,
) → torch.Tensor#

Load the hidden-state tensor for a single JSONL record.

Parameters:

  • record – A parsed JSONL line that contains hs_offset, hs_bytes, hs_shape.

  • trace_dir – Directory containing hidden_states_rank{rank}.bin.

Returns:

Tensor of shape [num_tokens, hidden_size] in bfloat16.

core.transformer.moe.router_trace.load_logits_for_record(record: dict, trace_dir: str) → torch.Tensor#

Load the pre-topk routing logits for a single JSONL record.

Parameters:

  • record – A parsed JSONL line that contains logit_offset, logit_bytes, logit_shape.

  • trace_dir – Directory containing logits_rank{rank}.bin.

Returns:

Tensor of shape [num_tokens, num_experts] in bfloat16.

class core.transformer.moe.router_trace.RouterTracer(
output_dir: str,
max_steps: int,
rank: int,
training_mode: bool = False,
capture_hidden_states: bool = False,
capture_logits: bool = False,
dump_router_weights: bool = False,
)#

Captures router top-K decisions across all MoE layers per step.

  • Inference mode: step boundaries are auto-detected. When a layer that has already fired this step fires again, a new step has started.

  • Training mode: the training loop calls advance_step() at each iteration boundary.

Recording is skipped during CUDA graph capture since D2H copies inside a captured graph would record stale values on replay.

Initialization

register_hooks(model) → None#

Walk model and register forward hooks on every TopKRouter module. Accepts a single model or a list of model chunks.

remove_hooks() → None#

Remove all forward hooks registered by register_hooks().

advance_step(step_id: Optional[int] = None) → None#

Advance to the next step (training mode).

Call once per training iteration, after the forward-backward pass. Flushes accumulated records to disk and disables the tracer once max_steps is reached.

Parameters:

step_id – Authoritative step id for the records just captured (e.g. the training iteration). When provided, buffered records are stamped with it and the tracer adopts it, so traces stay aligned with the caller’s step numbering (e.g. across checkpoint resumes) rather than a private 0-based counter. When omitted, the tracer falls back to incrementing its own counter.

make_hook(module_name: str = '')#

Build a forward hook callable for a single TopKRouter module.

The module’s qualified name is parsed once to recover its (block, mtp_idx, layer) identity so decoder and MTP layers that share a layer_number are kept distinct.

_extract_hidden_state(inputs, expected_num_tokens)#

Return a 2-D [num_tokens, hidden_size] bfloat16 tensor from hook inputs, or None.

_make_index_record(
top_indices_cpu,
step,
block,
mtp_idx,
layer,
) → dict#

Assemble a JSONL record dict for one layer’s top-K indices.

_record(module, inputs, outputs, identity=None) → None#
record_indices(
indices,
step: Optional[int] = None,
layer_ids: Optional[List[int]] = None,
block: str = 'decoder',
mtp_idx: Optional[int] = None,
) → None#

Serialize already-captured top-K routing indices through the JSONL sink.

This is the entry point for the in-pipeline recorder (RouterReplay/RoutingMetadata). Instead of capturing indices with a forward hook, the caller hands over the indices the router pipeline recorded. This works under CUDA graphs, because the recorder copies into a static buffer rather than relying on a Python hook firing during replay.

Only the top-K indices are serialized here. The hidden-state / logit / weight sidecars remain hook-only.

Parameters:

  • indices – Either a single tensor of shape [num_tokens, num_layers, topk] (the layout RoutingMetadata.get_routing_indices() returns), or a list/tuple of per-layer tensors each shaped [num_tokens, topk].

  • step – Step id stamped on the emitted records. Defaults to the tracer’s current step_id (drive boundaries with advance_step()).

  • layer_ids – Layer numbers, one per layer in indices. Defaults to range(num_layers).

  • block – Block tag for the records (“decoder” or “mtp”).

  • mtp_idx – MTP head index.

_flush_records_to_disk() → None#
flush() → None#

Flush remaining records.