Diff Coverage

        if weight is not None:
            # Ignore positions may contain -100, which is not a valid class
            # index. Use a harmless index there; their gradients are masked
            # below and therefore do not contribute to the result.
            safe_target = torch.where(
                ignore_mask, target_flat, torch.zeros_like(target_flat)
            )
            sample_weights = weight[safe_target]
        else:
            sample_weights = None

        if reduction == "mean":

        if reduction == "none":
            grad_scale_expanded = grad_scale.unsqueeze(-1)
        else:
            grad_scale_expanded = grad_scale.reshape(1, 1)
        if sample_weights is not None:
            grad_scale_expanded = grad_scale_expanded * sample_weights.unsqueeze(-1)
        grad_input = softmax_local * grad_scale_expanded

        local_targets = torch.where(in_vocab_mask, target_flat - vocab_start, torch.zeros_like(target_flat))

        if in_vocab_mask.any():

        if in_vocab_mask.any():
            row_indices = torch.arange(batch_size, device=target.device, dtype=torch.long)

            if reduction == "none":
                grad_values = -grad_scale
                if sample_weights is not None:
                    grad_values = grad_values * sample_weights
            else:
                grad_values = -grad_scale.expand_as(target_flat)
                if sample_weights is not None:
                    grad_values = grad_values * sample_weights

            grad_input = grad_input.contiguous()
            grad_input[row_indices[in_vocab_mask], local_targets[in_vocab_mask]] += grad_values[in_vocab_mask]

        if mask is not None:
            with torch.no_grad():
                if mask.device != labels.device:
                    mask = mask.to(labels.device)
                labels = labels.masked_fill(mask.view(-1) == 0, self.ignore_index)

        if self.fp32_upcast:
            logits = logits.float()

            gate_up_proj = gate_up_proj.view(self.num_local_experts, self.hidden_size, -1)

        # up
        if permuted.nelement() != 0:
            fc1_output = grouped_matmul(  # pylint: disable=not-callable
                permuted, gate_up_proj, bias=None, group_list=self._group_list, group_type=0, group_list_type=0,
            )
            if self.add_bias:
                b1 = self.bias1.view(self.num_local_experts, -1)
                fc1_output = fc1_output + torch.repeat_interleave(b1, self._tokens_per_expert_gmm, dim=0)
        else:
            gate_up_proj_2d = gate_up_proj.view(self.hidden_size, -1)
            fc1_output = torch.matmul(permuted, gate_up_proj_2d)

        if self.use_2d_experts:
            down_proj = down_proj.view(self.num_local_experts, -1, self.hidden_size)

        if fc1_output.nelement() != 0:
            fc2_output = grouped_matmul(  # pylint: disable=not-callable
                fc1_output, down_proj, bias=None, group_list=self._group_list, group_type=0, group_list_type=0,
            )
            if self.add_bias:
                b2 = self.bias2.view(self.num_local_experts, -1)
                fc2_output = fc2_output + torch.repeat_interleave(
                    b2, self._tokens_per_expert_gmm, dim=0,
                )
        else:

    ) -> torch.Tensor:
        """Run grouped experts with the Transformers Experts interface."""
        hidden_shape = hidden_states.shape
        hidden_states_flat = hidden_states.view(-1, hidden_states.shape[-1])
        permuted_tokens, sorted_indices = moe_token_permute(  # pylint: disable=not-callable
            hidden_states_flat,
            top_k_index,
        )
        flatten_probs = top_k_weights.view(-1)

            permuted_tokens,
            tokens_per_expert,
            permuted_probs,
        )
        output = moe_token_unpermute(  # pylint: disable=not-callable
            expert_outputs,
            sorted_indices,
            top_k_weights,
        )

        )

    def close(self) -> None:
        """Release the auxiliary process group after a dataset build."""
        if self._gloo_group is not None and dist.is_initialized():
            dist.destroy_process_group(self._gloo_group)
        self._gloo_group = None

    if builds_cache_first and owns_dataset:
        local_dataset = dataset_factory()

    if barrier_needed:
        try:
            dataloader_context.barrier()
        finally:
            # Dataset-specific barriers may own a temporary auxiliary group.
            # Close it after all ranks have crossed the barrier so repeated
            # dataset builds do not retain process groups and file handles.
            barrier = dataloader_context.barrier
            barrier_owner = getattr(barrier, "__self__", None)
            if barrier_owner is None and hasattr(barrier, "close"):
                barrier_owner = barrier
            close = getattr(barrier_owner, "close", None)
            if close is not None:
                close()

    if not builds_cache_first and owns_dataset:
        local_dataset = dataset_factory()

        print("Time to startup:", startup_end - startup_start)
        pack_to_seq_len = getattr(self.args, "pack_to_seq_len", None)
        chunk_size = pack_to_seq_len + 1 if pack_to_seq_len is not None else None
        token_buffers = {key: [] for key in keys}
        try:
            with open(input_file_name, "r", encoding="utf-8") as fin:
                encoded_docs = pool.imap(encoder.encode, fin, 32)
                for i, (doc, sentence_lens, bytes_processed) in enumerate(encoded_docs, start=1):
                    if self.args.find_optimal_num_workers and i > self.args.max_documents:
                        break
                    total_bytes_processed += bytes_processed
                    for key in keys:
                        if chunk_size is None:
                            builders[key].add_document(doc[key], sentence_lens[key])
                            continue
                        token_buffers[key].extend(doc[key])
                        complete_length = len(token_buffers[key]) // chunk_size * chunk_size
                        for offset in range(0, complete_length, chunk_size):
                            chunk = token_buffers[key][offset : offset + chunk_size]
                            builders[key].add_document(chunk, [chunk_size])
                        del token_buffers[key][:complete_length]
                    self.print_processing_stats(i, proc_start, total_bytes_processed)

            for key in keys:
                builders[key].finalize(output_idx_files[key])
        finally:
            for builder in builders.values():
                builder.data_file.close()
            pool.close()
            pool.join()

        return self.performance

    try:
        primary, secondaries = _classify_rewrite_result(
            apply_fn(), target, name)
        for request in secondaries:
            committed_secondaries.append(
                (request.target, _commit_forward_rewrite(request))
            )
        if primary is None:
            # In-place contract (external @inner_wrapper wrappers). Detect

    except Exception:
        # Failure rollback: restore every attribute written during the
        # rewrite so a half-installed wrapper never survives.
        for secondary_target, secondary_state in reversed(committed_secondaries):
            for attr, saved in secondary_state.items():
                _restore_attr(secondary_target, attr, saved)
        for attr, saved in saved_state.items():
            _restore_attr(target, attr, saved)
        raise
    target_name = _inner_target_name(module, target)

def _cp_sdpa_call(orig_sdpa, cp_mesh, q, k, v, kwargs, *, normalize_gqa=False):
    """CP-aware SDPA: K/V all-gather + D-04 offset-aware causal mask."""
    if normalize_gqa:
        q, k, v, kwargs = _normalize_hf_sdpa_gqa(q, k, v, kwargs)
    cp_dim = 2  # sequence dim of the [B, N, S, H] layout
    global_k, global_v = flex_cp_allgather(
        k.contiguous(), v.contiguous(), cp_dim, cp_mesh)
    if cp_mesh.size() > 1:

        def cp_aware_sdpa(q: Any, k: Any, v: Any, **kw: Any) -> Any:
            """CP-aware SDPA replacement: all-gather K/V plus the D-04 mask."""
            fired["hit"] = True
            return _cp_sdpa_call(
                orig_sdpa, cp_mesh, q, k, v, kw, normalize_gqa=True
            )

        F.scaled_dot_product_attention = cp_aware_sdpa

        self.profiler.start()

    def on_step_end(self, state: TrainerState, **kwargs: Any) -> None:
        """Advance the profiler schedule after one complete optimizer step."""
        del kwargs
        if self.profiler is not None:
            self.profiler.step()
            if state.global_step >= self.config.end_step:
                self.profiler.stop()
                self.profiler = None

    def on_train_end(self, state: TrainerState, **kwargs: Any) -> None:
        """Stop the profiler and flush any pending trace output."""
        del state, kwargs

        from hyper_parallel.core.dtensor.dtensor import _LAYOUT_CACHE  # pylint: disable=C0415
        from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution  # pylint: disable=C0415
        from hyper_parallel.core.fully_shard.hsdp_param import _GROUP_INFO_CACHE  # pylint: disable=C0415
        from hyper_parallel.platform.platform import EXISTING_COMM_GROUPS  # pylint: disable=C0415
        from hyper_parallel.platform.torch.platform import (  # pylint: disable=C0415
            _P2P_MULTI_STREAM_GROUPS,
        )

        EXISTING_COMM_GROUPS.clear()
        _P2P_MULTI_STREAM_GROUPS.clear()
        _DEVICE_MESH_MAP.clear()
        _LAYOUT_CACHE.clear()
        _GROUP_INFO_CACHE.clear()
        _HYBRID_MESH_CACHE.clear()

    # delegate ctx-manager behaviour
    def __enter__(self) -> Any:
        """Enter the wrapped profiler's context manager."""
        self.start()
        return self

    def __exit__(self, *a: Any) -> Any:
        """Exit the wrapped profiler's context manager."""
        del a
        self.stop()
        return False

    def start(self) -> Any:
        """Start profiling and begin recording the allocator history."""
        out = self._p.start()

        return out

    def stop(self) -> Any:
        """Stop profiling and stop recording the allocator history."""
        try:
            return self._p.stop()
        finally:
            get_torch_device().memory._record_memory_history(enabled=None)

    def step(self, *a: Any, **kw: Any) -> Any:
        """Advance the wrapped profiler by one step."""
        return self._p.step(*a, **kw)

        total_loss = 0.0
        total_loss_dict = defaultdict(int)

        for micro_step, batch in enumerate(training_batches):
            model_inputs, loss_inputs = batch
            self.base.model_reshard(micro_step, num_micro_steps)
            self.base.configure_fsdp_gradient_sync(
                micro_step,
                num_micro_steps,

            loss, loss_dict = self.base.forward_backward_step(model_inputs)

            # Release each device batch as soon as its backward pass completes;
            # prefetching all micro-batches must not pin them until step end.
            training_batches[micro_step] = None

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