Diff Coverage

    ctx.dp_rank = dist.get_rank() // tp
    return ctx


def train_fn(model, *, input_ids, labels):
    """Standard CE loss on the boundary-wrapped model forward.

    In SP mode the lm_head boundary all-gathers hidden_states to full
    sequence, so logits and labels both have the full seq_len.  The

            "[SP] sample input full seq_len=%s (SP sharding via embedding "
            "reduce-scatter)",
            input_batch.shape[1],
        )
    trainer.compile(input_ids=input_batch, labels=label_batch)
    inspect_graph(trainer)

    # 7. Training data iterator (DataSampler yields full-sequence batches)
    _LOG.info("\nStarting training...")

        """
        placeholders = [n for n in graph_module.graph.nodes if n.op == "placeholder"]
        num_user_inputs = len(placeholders) - num_state_inputs
        if num_user_inputs != 2:
            raise ValueError(
                f"PP v1 requires exactly two model inputs (input, label); "
                f"the traced graph carries {num_user_inputs}. Pass exactly "
                f"the input and label keyword arguments, or run without PP "
                f"(kwargs-general PP input routing is a follow-up)."

    user_flat, user_spec = torch.utils._pytree.tree_flatten(inputs)
    traced_spec = getattr(joint_graph.graph_module, "user_inputs_spec", None)
    if traced_spec is not None and user_spec != traced_spec:
        raise ValueError(
            "model inputs have a different pytree structure than during "
            "tracing (keys, nesting, or leaf types changed).\n"
            f"  Traced spec: {traced_spec}\n"
            f"  Got spec:    {user_spec}"

    def _place_on_device(self, inputs: dict) -> dict:
        """Move a dict of model inputs onto the compiler's device."""
        device = self._compiler.device
        if device is None:
            return inputs
        return {
            key: value.to(device) if isinstance(value, torch.Tensor) else value
            for key, value in inputs.items()
        }