refactor(cp): sharder ownership

[akoumpa]

What does this PR do?

Simplifies the context-parallel forward API introduced by #2937 and makes ContextParallelSharder the single recipe-facing owner of CP batch preparation and token layout.

The sharder now:

• derives TE and Magi behavior from the live model instead of accepting use_te, magi, or backend arguments;
• derives packed THD input from batch["qkv_format"];
• reuses a model-owned prepare_model_inputs_for_cp sharder when one is available;
• binds the CP/TP meshes once, so token shard/gather calls do not accept a mesh argument;
• shards the batch through shard(batch) and returns (train_ctx, sharded_batch);
• lives under nemo_automodel.components.distributed.context_parallel, with only ContextParallelSharder exported at the package level.

API

from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder

cp_sharder = ContextParallelSharder(
    model,
    device_mesh,
    batch,
    padding_token_id=tokenizer.pad_token_id,
)
train_ctx, batch = cp_sharder.shard(batch)

# Optional token-aligned tensors use the same captured CP layout.
local_advantages = cp_sharder.shard_token_tensor(
    advantages,
    fill=0.0,
)

with train_ctx():
    outputs = model(**batch)
    local_logprobs = compute_token_logprobs(outputs, batch["labels"])

full_logprobs = cp_sharder.gather_token_tensor(
    local_logprobs,
    trim=True,
)

There is intentionally no separate .context() call. shard(batch) performs the batch transformation, captures the resulting token layout, and returns the context factory required by the model forward.

shard_token_tensor(...) and gather_token_tensor(...) use the mesh and layout already bound to the sharder. This keeps auxiliary token tensors aligned with round-robin SDPA, contiguous model-owned layouts, packed THD partitions, and Magi dispatch permutations without adding them to the model batch.

A DTensor is not returned because standard Shard(dim) placements cannot represent all supported CP layouts, including head/tail ordering, packed THD partitions, Magi permutations, and model-owned re-padding.

Changelog

• Remove the recipe-facing prepare_cp_forward wrapper.
• Rename ContextParallelismSharder to ContextParallelSharder.
• Move CP implementation modules under components/distributed/context_parallel/.
• Update LLM, VLM, dLLM, KD, model-owned CP hooks, tests, and tools to use the new package and API.
• Keep backend-specific and model-specific preparation behind the sharder boundary.

Validation

• 138 passed across the focused CP sharder, CP utility, differential coverage, input-embedding, and Mamba CP unit tests.
• 142 passed, 1 skipped across Magi and model-owned CP tests.
• Ruff checks and syntax parsing passed.
• Verified that the previous root-level CP module import paths are removed.

Local recipe collection is currently blocked before reaching this code by the environment's incompatible mlflow and system cachetools installations.

Additional Information

Related to #2861 and #2937.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[akoumpa]

/ok to test 3197656

[HuiyingLi]

Thanks @akoumpa . We have a init and also a _from_strategy and then bind. I'm wondering whats the use case/if necessary?

[HuiyingLi]

Not super sure about using init sharder+sharder.shard(). The sharder init becomes a bit confusing. Could you please explain your consideration for replacing the function call? many thanks

[akoumpa]

Thanks @HuiyingLi, the goal was to make ContextParallelSharder the single owner of CP strategy and mesh state, rather than keeping a procedural wrapper around it.

The object-based API has two advantages:

1. The sharder can be cached on the trainer and reused, avoiding repeated strategy and mesh resolution. Only the batch-dependent ShardLayout is refreshed by each shard() call.
2. The sharder owns the device mesh and related options, so shard(), shard_token_tensor(), and gather_token_tensor() do not need the mesh and other runtime arguments passed on every call.

The separate .shard(batch) operation is intentional because that is where the batch is mutated and any data-dependent layout is computed and retained for subsequent token-aligned sharding or gathering.

[akoumpa]

/ok to test 47925da

[akoumpa]

/ok to test 6ccbc76