@@ -2,7 +2,7 @@
 
 import logging
 import warnings
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Tuple
 
 import torch
 
@@ -77,8 +77,11 @@ def __init__(self, mimo_config: MimoModelConfig, cp_group=None, tp_group=None) -
         max_seq_len = mimo_config.language_model_spec.params.get('max_sequence_length', 4096)
 
         self.partition_adapter: Optional[PartitionAdapter] = None
-        # Create partition adapter only if parallelism is enabled
-        if language_config.context_parallel_size > 1 or language_config.sequence_parallel:
+        # Only on language-module ranks: encoder-only ranks never shard and would read
+        # process groups they do not own.
+        if self.role.has_language_module and (
+            language_config.context_parallel_size > 1 or language_config.sequence_parallel
+        ):
             partition_config = PartitionConfig.from_mp_config(
                 mp=language_config,
                 max_seq_len=max_seq_len,
@@ -307,10 +310,24 @@ def get_text_embeddings(
         else:
             position_ids_text = position_ids[batch_idx, seq_idx].unsqueeze(0)
 
-        text_embeddings = (
-            unwrap_model(self.language_model)
-            .embedding(input_ids=input_ids_text, position_ids=position_ids_text)
-            .squeeze(1)
+        embedding_layer = unwrap_model(self.language_model).embedding
+        # Combined embeddings are SP-scattered later in PartitionAdapter; a second scatter
+        # here would split the flat text tokens across TP ranks before alignment.
+        if (
+            self.partition_adapter is not None
+            and self.partition_adapter.cfg.seq_parallel
+            and getattr(embedding_layer, 'scatter_to_sequence_parallel', False)
+        ):
+            raise RuntimeError(
+                "MIMO sequence parallelism requires the language embedding scatter to be "
+                "disabled; pass scatter_embedding_sequence_parallel=False when constructing "
+                "the language model."
+            )
+
+        text_embeddings = embedding_layer(
+            input_ids=input_ids_text, position_ids=position_ids_text
+        ).squeeze(
+            1
         )  # Shape: [num_text_tokens, hidden_dim]
         return text_embeddings
 
@@ -382,11 +399,14 @@ def forward(
                 return self._forward_encoders(input_ids, modality_inputs, input_tensors), loss_mask
 
             if self.role.has_language_module:
-                return (
-                    self._forward_language_module(
-                        input_ids, position_ids, attention_mask, labels, input_tensors
-                    ),
+                return self._forward_language_module(
+                    input_ids,
+                    position_ids,
+                    attention_mask,
                     loss_mask,
+                    labels,
+                    input_tensors,
+                    packing_kwargs,
                 )
 
             raise RuntimeError(f"Rank has no modules assigned in role: {self.role}")
@@ -497,28 +517,86 @@ def _empty_encoder_output(self, encoder_name: str) -> torch.Tensor:
             requires_grad=True,
         )
 
+    def _build_packed_seq_params(self, packing_kwargs: Optional[dict]) -> Optional[PackedSeqParams]:
+        """Build THD ``PackedSeqParams`` from ``packing_kwargs`` (None if not packing)."""
+        if packing_kwargs is None:
+            return None
+        for key in packing_kwargs:
+            if 'cu_seqlens' in key and packing_kwargs[key] is not None:
+                packing_kwargs[key] = packing_kwargs[key].to(dtype=torch.int32)
+        packed_seq_params = PackedSeqParams(**packing_kwargs)
+        packed_seq_params.qkv_format = 'thd'
+        return packed_seq_params
+
+    def _shard_language_inputs(
+        self,
+        embeddings: Optional[torch.Tensor],
+        labels: Optional[torch.Tensor],
+        loss_mask: Optional[torch.Tensor],
+        packed_seq_params: Optional[PackedSeqParams] = None,
+    ) -> Tuple[
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        Optional[PackedSeqParams],
+    ]:
+        """Apply CP/SP sharding via the partition adapter, or pass through if inactive.
+
+        ``embeddings`` are sequence-first ``(S, B, H)`` (``None`` on non-first PP stages)
+        and come back in ``(S/(cp*tp), B, H)``; labels/loss_mask are ``(B, S)``.
+        """
+        if self.partition_adapter is None:
+            return embeddings, labels, loss_mask, packed_seq_params
+
+        return self.partition_adapter.shard(
+            embeddings=embeddings,
+            labels=labels,
+            loss_mask=loss_mask,
+            packed_seq_params=packed_seq_params,
+        )
+
     def _forward_language_module(
         self,
         input_ids: torch.Tensor,
         position_ids: Optional[torch.Tensor],
         attention_mask: Optional[torch.Tensor],
+        loss_mask: Optional[torch.Tensor],
         labels: Optional[torch.Tensor],
         input_tensors: Optional[Dict[str, torch.Tensor]],
-    ) -> torch.Tensor:
+        packing_kwargs: Optional[dict] = None,
+    ) -> Tuple[Any, Optional[torch.Tensor]]:
         """Forward pass for language module on this rank.
 
         Args:
             input_ids: Token IDs
             position_ids: Position IDs
-            attention_mask: Attention mask
+            attention_mask: Attention mask. Must be ``None`` under context parallelism
+                (CP-local hidden states cannot line up with a dense mask); mask via a
+                causal ``attn_mask_type`` or ``packed_seq_params`` instead.
+            loss_mask: Loss mask for per-token loss normalization
             labels: Labels for loss computation
             input_tensors: Hidden states or embeddings from previous stage
+            packing_kwargs: Optional kwargs to construct packed (THD) sequence params.
 
         Returns:
-            Language model output (hidden states, logits, or loss depending on stage)
+            Tuple of (language model output, possibly CP-sharded loss mask). The
+            output is hidden states, logits, or loss depending on the stage.
         """
         lang_name = MIMO_LANGUAGE_MODULE_KEY
 
+        if (
+            self.partition_adapter is not None
+            and self.partition_adapter.cfg.use_cp
+            and attention_mask is not None
+        ):
+            raise RuntimeError(
+                "MIMO context parallelism requires attention_mask=None; mask via a causal "
+                "attn_mask_type or packed_seq_params (a dense mask cannot line up with the "
+                "CP-sharded sequence)."
+            )
+
+        packed_seq_params = self._build_packed_seq_params(packing_kwargs)
+
         if self.role.is_first_stage(lang_name):
             # First stage: receive encoder embeddings, combine with text, pass to LM
             # Build modality embeddings dict from encoder outputs
@@ -534,13 +612,21 @@ def _forward_language_module(
             )
             modality_embeddings["text"] = text_embeddings
 
-            # Combine all embeddings
+            # Combine all embeddings ([S, B, H])
             combined_embeddings = self.align_embeddings_by_token_positions(
                 modality_embeddings=modality_embeddings,
                 input_ids=input_ids,
                 special_token_ids=self.special_token_ids,
             )
 
+            # Apply CP/SP sharding; combined_embeddings returns in [S/(cp*tp), B, H].
+            combined_embeddings, labels, loss_mask, packed_seq_params = self._shard_language_inputs(
+                embeddings=combined_embeddings,
+                labels=labels,
+                loss_mask=loss_mask,
+                packed_seq_params=packed_seq_params,
+            )
+
             lm_output = self.language_model(
                 # decoder_input replaces the embedding lookup, so input_ids is
                 # unused here; position_ids is still consumed by mRoPE in models
@@ -550,9 +636,19 @@ def _forward_language_module(
                 decoder_input=combined_embeddings,
                 labels=labels,
                 attention_mask=attention_mask,
+                packed_seq_params=packed_seq_params,
             )
         else:
-            # Non-first stage: receive hidden states from previous LM stage
+            # Non-first stage: receive hidden states from previous LM stage.
+            # Labels/loss_mask still need CP sharding so the loss on the last stage
+            # lines up with the CP-local hidden states.
+            _, labels, loss_mask, packed_seq_params = self._shard_language_inputs(
+                embeddings=None,
+                labels=labels,
+                loss_mask=loss_mask,
+                packed_seq_params=packed_seq_params,
+            )
+
             hidden_states = input_tensors.get(lang_name) if input_tensors else None
 
             # Set input tensor on language model for PP (unwrap DDP to reach GPTModel)
@@ -569,13 +665,14 @@ def _forward_language_module(
                 decoder_input=None,
                 labels=labels,
                 attention_mask=attention_mask,
+                packed_seq_params=packed_seq_params,
             )
 
         # Key output for non-last stages so schedule can route to next LM stage
         if not self.role.is_last_stage(lang_name):
-            return {lang_name: lm_output}
+            return {lang_name: lm_output}, loss_mask
 
-        return lm_output
+        return lm_output, loss_mask
 
     def _build_colocated_communicators(self):
         grid_map = self.mimo_config.module_to_grid_map
@@ -632,16 +729,7 @@ def _forward_all_modules(
 
         This is the original behavior, preserved for backward compatibility.
         """
-        # If packing_kwargs is provided, construct PackedSeqParams
-        packed_seq_params = None
-        if packing_kwargs is not None:
-            # Ensure correct dtype for seqlens tensors
-            for key in packing_kwargs:
-                if 'cu_seqlens' in key and packing_kwargs[key] is not None:
-                    packing_kwargs[key] = packing_kwargs[key].to(dtype=torch.int32)
-            packed_seq_params = PackedSeqParams(**packing_kwargs)
-            packed_seq_params.qkv_format = 'thd'
-            logger.debug(f"Packed sequence parameters: {packed_seq_params}")
+        packed_seq_params = self._build_packed_seq_params(packing_kwargs)
 
         # 1. Process each modality to get embeddings
         modality_embeddings = {}
@@ -679,25 +767,14 @@ def _forward_all_modules(
         )
         logger.debug(f"Combined embeddings shape: {combined_embeddings.shape}")
 
-        # 3. If sharding is needed, apply PartitionAdapter.
-        # combined_embeddings is [S, B, H]; transpose to [B, S, H] for shard() which expects
-        # batch-first layout (required by get_batch_on_this_cp_rank). After CP sharding each
-        # rank holds [B, S/cp, H]; transpose back to [S/cp, B, H] for the language model.
-        if self.partition_adapter is not None:
-            combined_embeddings = combined_embeddings.transpose(0, 1).contiguous()  # [B, S, H]
-            combined_embeddings, labels, loss_mask, _, packed_seq_params = (
-                self.partition_adapter.shard(
-                    embeddings=combined_embeddings,
-                    labels=labels,
-                    loss_mask=loss_mask,
-                    attention_mask=attention_mask,
-                    packed_seq_params=packed_seq_params,
-                )
-            )
-            # shard() returns embeddings in [B, S/cp, H]; transpose to [S/cp, B, H]
-            # which is what the language model expects.
-            if combined_embeddings is not None:
-                combined_embeddings = combined_embeddings.transpose(0, 1).contiguous()
+        # 3. Apply CP/SP sharding. combined_embeddings is [S, B, H] and returns
+        # [S/(cp*tp), B, H] for the LM (the adapter handles the CP batch-first transpose).
+        combined_embeddings, labels, loss_mask, packed_seq_params = self._shard_language_inputs(
+            embeddings=combined_embeddings,
+            labels=labels,
+            loss_mask=loss_mask,
+            packed_seq_params=packed_seq_params,
+        )
 
         # 5. Forward pass through language model
         lm_output = self.language_model(