fix(cp): preserve gradients when sharding VLM inputs

nemo_automodel/components/distributed/cp_utils.py

@@ -101,6 +101,21 @@ def create_context_parallel_ctx(
     )
 
 
+def _shard_grad_buffer_for_cp(buffer: torch.Tensor, seq_dim: int, cp_mesh: DeviceMesh) -> torch.Tensor:
+    """Shard a gradient-bearing buffer with CP's head-tail load-balancing order."""
+    cp_size = cp_mesh.size()
+    num_chunks = 2 * cp_size
+    seq_len = buffer.shape[seq_dim]
+    if seq_len % num_chunks != 0:
+        raise ValueError(f"CP sequence length {seq_len} must be divisible by {num_chunks}")
+
+    chunk_size = seq_len // num_chunks
+    cp_rank = cp_mesh.get_local_rank()
+    head_chunk = buffer.narrow(seq_dim, cp_rank * chunk_size, chunk_size)
+    tail_chunk = buffer.narrow(seq_dim, (num_chunks - cp_rank - 1) * chunk_size, chunk_size)
+    return torch.cat((head_chunk, tail_chunk), dim=seq_dim)
+
+
 def make_target_cp_ctx(cp_mesh: DeviceMesh, input_ids, position_ids=None):
     """Build a context-parallel context for a frozen target forward.
 
@@ -443,6 +458,19 @@ def _get_mesh_size(mesh):
         for idx, key in batch_buffer_keys.items():
             batch[key] = cp_buffers[idx]
 
+    # PyTorch's legacy context_parallel buffers API shards in place with
+    # ``resize_``/``copy_``. ``resize_`` rejects tensors that require gradients,
+    # and detaching inputs_embeds here would silently stop gradients to trainable
+    # embeddings and multimodal towers. Apply the same default head-tail shard
+    # out of place so autograd remains connected, then let context_parallel
+    # mutate only the integer/mask buffers.
+    primary_seq_tensor = cp_buffers[0]
+    if primary_seq_tensor.requires_grad:
+        batch[primary_key] = _shard_grad_buffer_for_cp(primary_seq_tensor, cp_seq_dims[0], cp_mesh)
+        cp_no_restore_buffers.remove(primary_seq_tensor)
+        cp_buffers = cp_buffers[1:]
+        cp_seq_dims = cp_seq_dims[1:]
+
     cp_ctx = create_context_parallel_ctx(
         cp_mesh=cp_mesh,
         cp_buffers=cp_buffers,

tests/unit_tests/distributed/test_cp_utils_inputs_embeds.py

@@ -33,20 +33,24 @@
 
 
 class _DummySubMesh:
-    def __init__(self, size: int):
+    def __init__(self, size: int, local_rank: int = 0):
         self._size = size
+        self._local_rank = local_rank
 
     def size(self) -> int:
         return self._size
 
     def get_group(self):
         return None
 
+    def get_local_rank(self) -> int:
+        return self._local_rank
+
 
 class _DummyDeviceMesh(dict):
-    def __init__(self, cp_size: int, tp_size: int):
+    def __init__(self, cp_size: int, tp_size: int, cp_rank: int = 0):
         super().__init__()
-        self["cp"] = _DummySubMesh(cp_size)
+        self["cp"] = _DummySubMesh(cp_size, cp_rank)
         self["tp"] = _DummySubMesh(tp_size)
         self.mesh_dim_names = ["cp", "tp"]
 
@@ -98,6 +102,34 @@ def _fake_create_ctx(**kwargs):
     assert cp_buffers[1] is labels
 
 
+def test_inputs_embeds_with_grad_is_sharded_out_of_place(monkeypatch):
+    """Grad-bearing embeds must bypass context_parallel's in-place resize."""
+    captured = {}
+
+    def _fake_create_ctx(**kwargs):
+        captured.update(kwargs)
+        return object()
+
+    monkeypatch.setattr(_cu, "create_context_parallel_ctx", _fake_create_ctx)
+    monkeypatch.setattr(_cu, "get_train_context", lambda *a, **kw: "ctx")
+
+    device_mesh = _DummyDeviceMesh(cp_size=2, tp_size=1, cp_rank=1)
+    full_sequence = torch.arange(8.0, requires_grad=True)
+    inputs_embeds = full_sequence.view(1, 8, 1)
+    batch = {
+        "inputs_embeds": inputs_embeds,
+        "labels": torch.arange(8).view(1, 8),
+    }
+
+    _cu.make_cp_batch_and_ctx(device_mesh, batch)
+
+    assert all(not buffer.requires_grad for buffer in captured["cp_buffers"])
+    assert torch.equal(batch["inputs_embeds"].flatten(), torch.tensor([2.0, 3.0, 4.0, 5.0]))
+
+    batch["inputs_embeds"].sum().backward()
+    assert torch.equal(full_sequence.grad, torch.tensor([0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0]))
+
+
 def test_input_ids_path_unchanged(monkeypatch):
     """Standard LLM path (input_ids only) unchanged by the inputs_embeds extension."""
     captured = {}