Diff Coverage

            if num_loss_outputs
            else list(outputs[1:])
        )
    else:
        loss, grads, loss_dict = outputs, [], {}
    return loss, grads, loss_dict

    ``TracingContext.get()`` raises ``RuntimeError`` when no trace is active,
    so we catch it and return False.
    """
    if torch.compiler.is_compiling():
        return True
    try:
        return torch._guards.TracingContext.get() is not None  # pylint: disable=protected-access
    except Exception:  # pylint: disable=W0718
        return False

def _reduce_op_to_str(reduce_op: Any) -> str:
    """Map a ``dist.ReduceOp`` to the string name ``_functional_collectives`` expects."""
    op_str = {
        dist.ReduceOp.SUM: 'sum',
        dist.ReduceOp.PRODUCT: 'product',
        dist.ReduceOp.MAX: 'max',
        dist.ReduceOp.MIN: 'min',

        dist.ReduceOp.PRODUCT: 'product',
        dist.ReduceOp.MAX: 'max',
        dist.ReduceOp.MIN: 'min',
    }.get(reduce_op, 'sum')
    if hasattr(dist.ReduceOp, 'AVG') and reduce_op == dist.ReduceOp.AVG:
        op_str = 'avg'
    return op_str


class _TorchContiguousGrad(torch.autograd.Function):  # pylint: disable=abstract-method
    """Autograd identity that materializes gradients before upstream collectives."""

    if rank_list is not None:
        group_ranks = dist.get_process_group_ranks(group)
        if tuple(rank_list) != tuple(group_ranks):
            # Rank subset gathering (not common path) — same for both modes.
            output = [
                _TorchContiguousGrad.apply(tensor)
                for tensor in dist_func.all_gather(data, group=group)
            ]
            rank_to_idx = {int(rank): idx for idx, rank in enumerate(group_ranks)}

                for tensor in dist_func.all_gather(data, group=group)
            ]
            rank_to_idx = {int(rank): idx for idx, rank in enumerate(group_ranks)}
            output = [output[rank_to_idx[int(rank)]] for rank in rank_list]
            return torch.cat(output, dim=concat_dim)
    if not _is_tracing_or_compiling():
        # Eager: original implementation using torch.distributed.nn.functional.
        output = [
            _TorchContiguousGrad.apply(tensor)

    # Compile: use functional collective API for torch.compile compatibility.
    # all_gather_single only reliably supports gather_dim=0; for other dims we
    # gather on dim 0 then reshape to avoid _maybe_view_chunk_cat view() bugs
    # in Dynamo fake tensor tracing.
    group_size = dist.get_world_size(group)
    import torch.distributed._functional_collectives as fc  # pylint: disable=C0415
    gathered = fc.all_gather_single(data, 0, group)
    if hasattr(gathered, 'wait'):
        gathered = gathered.wait()
    if concat_dim == 0:
        return gathered
    # Reshape from [group_size * d0, d1, ...] to [d0 * group_size, d1, ...]
    # by moving the group_size factor from dim 0 to concat_dim.
    orig_shape = list(data.shape)
    gathered_shape = [orig_shape[0] * group_size] + orig_shape[1:]
    gathered = gathered.view(gathered_shape)
    # Move group_size from dim 0 to concat_dim via unflatten + movedim + flatten
    unflattened = gathered.unflatten(0, (group_size, orig_shape[0]))
    moved = torch.movedim(unflattened, 0, concat_dim)
    result_shape = list(orig_shape)
    result_shape[concat_dim] *= group_size
    return moved.reshape(result_shape).contiguous()


def differentiable_all_to_all_single(input_tensor: Tensor, input_splits: Sequence[int],
                                     output_splits: Sequence[int], group) -> Tensor:

    if not _is_tracing_or_compiling():
        return dist_func.all_reduce(data, op=reduce_op, group=group)
    # Compile: use functional collective API for torch.compile compatibility.
    # fc.all_reduce expects a string reduceOp and a ProcessGroup.
    import torch.distributed._functional_collectives as fc  # pylint: disable=C0415
    return fc.all_reduce(data, _reduce_op_to_str(reduce_op), group)


def differentiable_reduce_scatter(data: Tensor, dev_num: int, axis: int,
                                  op: Union[str, Any], group) -> Tensor:

            output_tensor, input_tuple, op=reduce_op, group=group
        )
        # 'avg' maps to SUM in _OP_MAP, so the division stays manual.
        if op == 'avg':
            output_tensor = output_tensor / dev_num
        return output_tensor

    # Compile: use functional collective API for torch.compile compatibility.
    # fc.reduce_scatter_single expects a string reduceOp, scatter_dim, and group.
    import torch.distributed._functional_collectives as fc  # pylint: disable=C0415
    result = fc.reduce_scatter_single(data, _reduce_op_to_str(reduce_op), axis, group)
    if hasattr(result, 'wait'):
        result = result.wait()
    # 'avg' maps to SUM in _OP_MAP, so the division stays manual.
    if op == 'avg':
        result = result / dev_num
    return result


def differentiable_variable_all_gather(
        input_tensor: Tensor, output_splits: Sequence[int], group) -> Tensor:

            self.model_fwd_context = partial(loss_parallel, mesh=tp_mesh)
        else:
            self.model_fwd_context = nullcontext()
        self.model_bwd_context = nullcontext()
        self.graph_compiler: Optional[GraphCompiler] = None

    def _build_graph_compiler(self) -> None:
        """Build the optional joint-graph compiler for this trainer runtime."""
        if not self.config.compile.selects_graph_compiler():

                    model_inputs=micro_batch,
                    labels=labels,
                )
            else:
                model_fwd_context = (
                    self.model_fwd_context()
                    if callable(self.model_fwd_context)
                    else self.model_fwd_context
                )
                with model_fwd_context:
                    outputs: ModelOutput = self.model(**micro_batch, use_cache=False)

                # with use_parallel_state("base"):
                loss, loss_dict = self.postforward(outputs, labels)
                # The loss graph owns everything required for backward. Releasing
                # the model output here avoids retaining large vocabulary logits
                # until the whole backward pass finishes.
                del outputs
                if self.config.training.empty_cache_before_backward:
                    empty_cache()

                # Backward pass
                with self.model_bwd_context:
                    loss.backward()

            del micro_batch
            return loss, loss_dict

        raise ConfigResolutionError(path, "DataLoader configuration must be a YAML mapping")

    target_node = dict(node)
    if "dataloader_type" in target_node:
        raise ConfigResolutionError(f"{path}.dataloader_type", "renamed to sampler_type")
    if "batch_adapter" in target_node:
        raise ConfigResolutionError(
            f"{path}.batch_adapter",
            "removed by the Omni data lifecycle; use dataset.data_transform for "
            "sample/batch encoding and dataloader.get_batch.runtime_input_adapter "
            "for model-owned forward metadata",

        loss_name = key.split("_loss", maxsplit=1)[0]  # foundation/image_decoder/**

        cur_token_len = current_token_counts[f"{loss_name}_tokens"]
        if sequence_parallel:
            if _is_tracing_or_compiling():
                # Graph tracing: ``.item()`` is data-dependent and unsupported
                # by the static-shape tracer; aggregate on tensors instead.
                cur_token_len = differentiable_all_reduce(
                    cur_token_len, "sum", sequence_parallel_group
                )
            else:
                cur_token_len = all_reduce(cur_token_len.item(), op="sum", group=sequence_parallel_group)

        if _is_tracing_or_compiling():
            # Graph-traced variant of the aggregation below: identical math
            # expressed with tensor collectives so the static-shape tracer can
            # capture it (no ``.item()``, no data-dependent zero-token guard).
            # ``differentiable_all_reduce`` dispatches to functional collectives

            # capture it (no ``.item()``, no data-dependent zero-token guard).
            # ``differentiable_all_reduce`` dispatches to functional collectives
            # under tracing; the token-count tensors are baked as trace-time
            # constants, matching the tracer's fixed-shape contract.
            step_len = differentiable_all_reduce(
                step_token_counts[f"{loss_name}_tokens"], "sum", dp_cp_group
            )
            local_weighted_loss = cur_loss * cur_token_len
            backward_loss = local_weighted_loss / step_len * device_mesh.dp_size * device_mesh.cp_size
            global_weighted_loss = differentiable_all_reduce(
                local_weighted_loss.detach(), "sum", dp_cp_group
            )
            global_mean = global_weighted_loss / step_len
            cur_loss = backward_loss + global_mean - backward_loss.detach()
            if sequence_parallel:
                cur_loss = cur_loss / sequence_parallel_size
            loss_dict[key] = cur_loss
            continue

        all_reduced_len = all_reduce(
            step_token_counts[f"{loss_name}_tokens"].item(),
            op="sum",

        self.base._build_optimizer()
        self.base.model_integration.attach_optimizer(self.base.optimizer)
        self.base._build_lr_scheduler()
        self.base._build_training_context()
        self.base._build_graph_compiler()
        self.base._init_callbacks()

    def _build_model_assets(self) -> None:
        """Build tokenizer-backed assets for text training."""

        # Checkpoint and logging callbacks observe the number of completed
        # optimizer updates.
        self.base.state.global_step += 1
        grad_norm_value = float(grad_norm)
        self.base._end_model_integration_step(  # pylint: disable=protected-access
            {"loss": total_loss, "grad_norm": grad_norm_value}
        )
        self.on_step_end(
            loss=total_loss,