Diff Coverage

            direct_backward: bool,
        ) -> tuple:
            """Compute and retain the accumulated chunk gradients."""
            if torch.is_tensor(loss_scale):
                if loss_scale.numel() != 1:
                    raise ValueError(
                        "loss_scale must be a scalar, "
                        f"got shape {tuple(loss_scale.shape)}"
                    )
                if loss_scale.requires_grad:
                    raise ValueError("loss_scale must not require gradients")
                chunk_loss_scale = loss_scale.to(
                    device=hidden_states.device, dtype=torch.float32,
                )
            else:
                chunk_loss_scale = float(loss_scale)

            """Replay the private input and output-head gradient buffers."""
            del grad_loss_sum
            grad_hidden, grad_weight = ctx.saved_tensors
            if grad_scaled_loss is None:
                return None, None, None, None, None, None, None

            # The Trainer passes the complete scale into ``forward`` and calls
            # backward directly on ``scaled_loss_sum``. Returning the private
            # buffers for that unit upstream gradient avoids an NPU scalar-mul

    Raises:
        ValueError: If an input shape or option is invalid.
    """
    if chunk_size <= 0:
        raise ValueError(f"chunk_size must be positive, got {chunk_size}")
    if hidden_states.dim() != 3:
        raise ValueError(
            "hidden_states must have shape [batch, sequence, hidden], "
            f"got {tuple(hidden_states.shape)}"
        )
    if targets.dim() != 2 or hidden_states.shape[:2] != targets.shape:
        raise ValueError(
            "targets must match hidden_states batch/sequence dimensions, "
            f"got hidden={tuple(hidden_states.shape)} and targets={tuple(targets.shape)}"
        )
    if hidden_states.size(1) == 0:
        raise ValueError("hidden_states sequence dimension must be non-empty")
    if head_weight.dim() != 2 or head_weight.size(1) != hidden_states.size(2):
        raise ValueError(
            "head_weight must have shape [vocab, hidden] matching hidden_states, "
            f"got hidden={tuple(hidden_states.shape)} and weight={tuple(head_weight.shape)}"
        )

        return self._output_loss_implementation() == "chunk_loss"

    def _validate_output_loss_configuration(self) -> None:
        """Reject unsupported output-loss combinations before model build."""
        implementation = self._output_loss_implementation()
        valid_implementations = {"eager", "chunk_loss"}
        if implementation not in valid_implementations:
            raise ValueError(
                "train.output_loss.implementation must be one of "
                f"{sorted(valid_implementations)}, got {implementation!r}."
            )
        if implementation == "eager":
            return
        tp_size = self._parallel_dim_size("tp")
        pp_size = self._parallel_dim_size("pp")
        if tp_size != 1 or pp_size != 1:
            raise ValueError(
                "Chunk Loss version one requires "
                "train.accelerator.tp=1 and pp=1, "
                f"got tp={tp_size}, pp={pp_size}."
            )
        chunk_size = int(self.args.train.output_loss.chunk_size)
        if chunk_size <= 0:
            raise ValueError(
                "train.output_loss.chunk_size must be positive for "
                f"chunk_loss, got {chunk_size}."
            )

        self.parallel_dims = ParallelDims.from_config(
            self.args.train.accelerator, world_size=platform.get_world_size(),
        )
        self._validate_output_loss_configuration()
        logger.info_rank0("ParallelDims: %s", self.parallel_dims.summary())
        # Mixed precision lives in FSDP2's MixedPrecisionPolicy, so a
        # low-precision run needs a dp_shard axis (size-1 is enough) for the
        # FSDP wrap to exist — see ``build_mesh``'s force_dp_shard contract.

        shifted_labels = micro_batch.pop("labels", None) if labels_are_shifted else None
        if labels_are_shifted and shifted_labels is None:
            raise ValueError("CP-shifted loss marker is set but labels are missing.")
        if labels_are_shifted and self._chunk_loss_enabled():
            micro_batch["loss_targets"] = shifted_labels
        return micro_batch, labels_are_shifted, shifted_labels

    def _compute_micro_loss(
        self,

        """Return mean loss and summed loss for one micro-batch."""
        if not labels_are_shifted:
            loss = outputs["loss"] if isinstance(outputs, dict) else outputs.loss
            if isinstance(outputs, dict) and outputs.get("loss_sum") is not None:
                return loss, outputs["loss_sum"].detach()
            return loss, loss.detach() * max(micro_batch_tokens, 1)

        if isinstance(outputs, dict) and outputs.get("connected_loss_sum") is not None:
            loss_sum = outputs["connected_loss_sum"]
            return loss_sum / max(micro_batch_tokens, 1), loss_sum

        if isinstance(outputs, dict) and outputs.get("loss_sum") is not None:
            loss_sum = outputs["loss_sum"]
            return loss_sum / max(micro_batch_tokens, 1), loss_sum

        logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits
        target_device = logits.device if hasattr(logits, "device") else self.device
        shifted_labels = shifted_labels.to(target_device, non_blocking=True)

        forward; otherwise backward would scale another full-size dWeight.
        """
        if not self._chunk_loss_enabled():
            return None
        if labels_are_shifted:
            if micro_batch.get("loss_targets") is None:
                return None
        elif micro_batch.get("labels") is None:
            return None

        def _model_loss_normalizer():
            num_items = micro_batch.get("num_items_in_batch")
            if num_items is not None:
                if torch.is_tensor(num_items):
                    return num_items.to(dtype=torch.float32).clamp_min(1)
                return float(max(num_items, 1))

            labels = micro_batch["labels"]
            targets = torch.nn.functional.pad(
                labels, (0, 1), value=-100,
            )[..., 1:]
            return targets.ne(-100).sum().to(dtype=torch.float32).clamp_min(1)

        agg = self.args.train.optimizer.loss_aggregation
        cp_size = self._cp_size()
        cp_rank_average = agg == "rank_average" and cp_size > 1
        if agg == "rank_average" and not cp_rank_average:
            normalizer = (
                float(max(micro_batch_tokens, 1))
                if labels_are_shifted
                else _model_loss_normalizer()
            )
            rank_scale = getattr(
                self.model, "hp_rank_average_loss_scale_size", 1,
            )
            loss_scale = 1.0 / normalizer
            if num_micro > 1:
                loss_scale = loss_scale / num_micro
            if rank_scale != 1:
                loss_scale = loss_scale / rank_scale
            return loss_scale

        loss_scale_size = getattr(
            self.model, "hp_token_loss_scale_size", self.parallel_dims.dp_size,
        )
        if labels_are_shifted:
            loss_scale = loss_scale_size / float(max(global_tokens, 1))
        else:
            normalizer = _model_loss_normalizer()
            loss_scale = (
                float(micro_batch_tokens)
                / float(max(global_tokens, 1))
                * loss_scale_size
                / normalizer

                * loss_scale_size
                / normalizer
            )

        tp_loss_scale_size = getattr(
            self.model,
            "hp_loss_tp_scale_size",
            max(1, self._parallel_dim_size("tp")),
        )
        if tp_loss_scale_size != 1:
            loss_scale = loss_scale / tp_loss_scale_size
        ep_loss_scale_size = getattr(self.model, "hp_loss_ep_scale_size", 1)
        if ep_loss_scale_size != 1:
            loss_scale = loss_scale / ep_loss_scale_size
        return loss_scale

    def _scale_loss_for_backward(
        self,
        loss,

            global_tokens,
            num_micro,
        )
        if output_loss_scale is not None:
            micro_batch["output_loss_scale"] = output_loss_scale

        # Forward (with training context for activation offload)
        with self.model_fwd_context:
            outputs = self.model(**micro_batch, use_cache=False)

            if isinstance(outputs, dict)
            else getattr(outputs, "backward_loss", None)
        )
        if backward_loss is not None:
            scaled_loss = backward_loss
        else:
            scaled_loss = self._scale_loss_for_backward(
                loss,
                loss_sum,