Diff Coverage

        The runtime calls this before capture. Eager retains its real boundaries;
        graph inputs carry token-shaped metadata, not a variable document list.
        """
        if self.local_query_start != 0 or self.local_query_length != self.global_sequence_length:
            return
        self.segment_start_mask(device)
        object.__setattr__(self, "_use_prepared_packing_inputs", True)

    def router_metadata(self, device: torch.device) -> tuple[torch.Tensor, int]:

        """Expose packing tensors with a schema independent of document count and S."""
        if self._use_prepared_packing_inputs:
            return [self._prepared_starts, self._prepared_positions,
                    self._prepared_masks, self._prepared_sequences], ("packing",)
        return [self.cu_seq_lens], (self.local_query_start, self.local_query_length, self.global_sequence_length)

    @classmethod
    def pytree_unflatten(cls, children: Any, context: tuple) -> SharedCompressedPackedSequence:
        """Restore prepared metadata without reading boundary values during capture."""

    @classmethod
    def pytree_unflatten(cls, children: Any, context: tuple) -> SharedCompressedPackedSequence:
        """Restore prepared metadata without reading boundary values during capture."""
        if context != ("packing",):
            return cls(children[0], *context)
        starts, positions, masks, sequences = children
        reference = next(iter(starts.values()))
        length = reference.shape[-1]
        packed = cls(reference.new_empty(0), 0, length, length)

        Returns:
            Inputs used for both graph capture and replay. Eager never calls this hook.
        """
        return model_inputs

    def runtime_input_fields(self) -> tuple[str, ...]:
        """Return the forward fields exclusively owned by this adapter.

                loss_dict[name] = value.detach().float()
        # Log-only outputs share the eager/graph path without entering backward loss.
        moe_monitor = getattr(self, "moe_monitor", None)
        if moe_monitor is not None:
            aux_loss = moe_monitor.last_mean_aux_loss
            if aux_loss is not None:
                loss_dict["router_aux_loss"] = aux_loss
        return loss, loss_dict

    def forward_backward_step(
            self,

            loss_dict: Dict[str, torch.Tensor]
            # token num for fixed_ce_loss in postforward
            self.current_token_counts = count_loss_token(micro_batch)
            loss, loss_dict = self.forward_backward_step(micro_batch)
            self.on_micro_step_end(loss_dict)

            total_loss += loss.item()
            for k, v in loss_dict.items():
                total_loss_dict[k] += v.item()

    def on_step_begin(
            self, state: TrainerState, micro_batches: list[dict[str, Any]] | None = None, **kwargs: Any,
    ) -> None:
        """Reset per-step metric state, including steps that do not emit a log."""
        del state, micro_batches, kwargs
        self._router_aux_sum = None
        self._router_aux_count = 0
        self._micro_step_outputs_seen = False

    def on_micro_step_end(self, state: TrainerState, loss_dict: dict[str, Any], **kwargs: Any) -> None:
        """Accumulate detached router outputs without synchronizing device scalars."""
        del state, kwargs
        if self.logging_steps <= 0 or getattr(self.trainer, "global_rank", 0) != 0:
            return
        self._micro_step_outputs_seen = True
        aux_loss = loss_dict.get("router_aux_loss") if loss_dict is not None else None
        if aux_loss is not None:
            value = aux_loss.detach()
            if self._router_aux_sum is None:
                # Graph executors may reuse output storage on the next microbatch.
                self._router_aux_sum = value.clone()
            else:
                self._router_aux_sum += value
            self._router_aux_count += 1

    @staticmethod
    def _format_value(value: Any) -> str:
        """Format a scalar metric compactly for one terminal line."""

        metrics = dict(getattr(self.trainer, "step_env_metrics", {}))
        moe_monitor = getattr(self.trainer, "moe_monitor", None)
        if self._micro_step_outputs_seen:
            aux_loss = (self._router_aux_sum / self._router_aux_count
                        if self._router_aux_count else None)
        else:
            # Keep direct step-end callers compatible when no micro-step outputs were observed.
            aux_loss = loss_dict.get("router_aux_loss") if loss_dict is not None else None