diff --git a/nemo_rl/models/automodel/data.py b/nemo_rl/models/automodel/data.py
index b272cc24d4..3ffbbc4d0a 100644
--- a/nemo_rl/models/automodel/data.py
+++ b/nemo_rl/models/automodel/data.py
@@ -16,7 +16,7 @@
 
 import itertools
 from dataclasses import dataclass, field
-from typing import Any, Iterator, Optional
+from typing import Any, Iterator, Optional, Tuple
 
 import torch
 from transformers import AutoTokenizer
@@ -340,3 +340,28 @@ def process_global_batch(
         "global_valid_seqs": global_valid_seqs,
         "global_valid_toks": global_valid_toks,
     }
+
+
+def check_sequence_dim(data: BatchedDataDict[Any]) -> Tuple[int, int]:
+    """Check and validate sequence dimension across all tensors.
+
+    Verifies that dimension 1 is the sequence dimension for all tensors
+    in the data dictionary that have more than one dimension.
+
+    Args:
+        data: BatchedDataDict to validate
+
+    Returns:
+        Tuple of (sequence_dim, seq_dim_size)
+
+    Raises:
+        AssertionError: If any tensor has inconsistent sequence dimension
+    """
+    sequence_dim = 1
+    seq_dim_size = data.get("input_ids").shape[sequence_dim]
+    for _, v in data.items():
+        if torch.is_tensor(v) and len(v.shape) > 1:
+            assert v.shape[sequence_dim] == seq_dim_size, (
+                f"Dim 1 must be the sequence dim, expected dim 1={seq_dim_size} but got shape {v.shape}"
+            )
+    return sequence_dim, seq_dim_size
diff --git a/nemo_rl/models/automodel/train.py b/nemo_rl/models/automodel/train.py
new file mode 100644
index 0000000000..d93a748942
--- /dev/null
+++ b/nemo_rl/models/automodel/train.py
@@ -0,0 +1,925 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Training utilities for automodel (DTensor-based) policy workers.
+
+This module provides post-processor classes and forward/backward functions
+that follow the same pattern as nemo_rl/models/megatron/train.py.
+
+Key differences from megatron approach:
+- Post-processors compute results directly (no callable return pattern)
+- forward_with_post_processing_fn calls post-processor directly
+- automodel_forward_backward uses PyTorch autograd instead of Megatron's pipeline
+"""
+
+from collections import defaultdict
+from typing import Any, Callable, Iterator, Optional, Tuple, Union
+
+import torch
+from nemo_automodel.components.distributed.tensor_utils import to_local_if_dtensor
+from torch import nn
+from torch.distributed.tensor import DTensor, Shard
+
+from nemo_rl.algorithms.interfaces import LossFunction
+from nemo_rl.algorithms.loss_functions import SequencePackingLossWrapper
+from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+from nemo_rl.distributed.model_utils import (
+    allgather_cp_sharded_tensor,
+    distributed_vocab_topk,
+    get_logprobs_from_vocab_parallel_logits,
+)
+from nemo_rl.models.automodel.data import ProcessedInputs, ProcessedMicrobatch
+from nemo_rl.models.policy import PolicyConfig
+
+# Union type for any post-processing function
+PostProcessingFunction = Union[
+    "LossPostProcessor",
+    "LogprobsPostProcessor",
+    "TopkLogitsPostProcessor",
+    "ScorePostProcessor",
+]
+
+
+def model_forward(
+    model: nn.Module,
+    processed_inputs: ProcessedInputs,
+    is_reward_model: bool = False,
+    allow_flash_attn_args: bool = True,
+) -> torch.Tensor:
+    """Perform a single forward pass through the model.
+
+    Args:
+        model: The model to run forward pass on
+        processed_inputs: ProcessedInputs containing all tensors for forward pass
+        is_reward_model: Whether this is a reward model
+        allow_flash_attn_args: Whether to pass flash_attn_kwargs to model
+
+    Returns:
+        torch.Tensor: Output tensor from the model (logits)
+    """
+    model_args = dict(
+        input_ids=processed_inputs.input_ids,
+        attention_mask=processed_inputs.attention_mask,
+        position_ids=processed_inputs.position_ids,
+        use_cache=False,
+    )
+
+    # Add flash attention kwargs if applicable
+    if processed_inputs.has_flash_attention:
+        model_args["flash_attn_kwargs"] = processed_inputs.flash_attn_kwargs
+
+    # Add VLM kwargs if applicable
+    if processed_inputs.is_multimodal:
+        model_args.update(processed_inputs.vlm_kwargs)
+        # flash_attn_kwargs is not supported for multimodal
+        if "flash_attn_kwargs" in model_args:
+            del model_args["flash_attn_kwargs"]
+
+    # Reward models don't support flash_attn_kwargs
+    if is_reward_model:
+        if "flash_attn_kwargs" in model_args:
+            del model_args["flash_attn_kwargs"]
+
+    # Remove flash_attn_kwargs if not allowed
+    if not allow_flash_attn_args and "flash_attn_kwargs" in model_args:
+        del model_args["flash_attn_kwargs"]
+
+    outputs = model(**model_args)
+    return outputs
+
+
+def extract_logits(
+    model: nn.Module,
+    outputs: Any,
+) -> torch.Tensor:
+    """Extract logits from model outputs.
+
+    Args:
+        model: The model (used for lm_head if needed)
+        outputs: Model outputs (can be tensor, DTensor, or object with logits attribute)
+
+    Returns:
+        torch.Tensor: Logits tensor
+    """
+    if isinstance(outputs, (torch.Tensor, DTensor)):
+        # Custom models can output logits directly
+        return outputs
+    elif not hasattr(outputs, "logits"):
+        return model.lm_head(outputs.last_hidden_state)
+    else:
+        return outputs.logits
+
+
+def apply_temperature_scaling(
+    logits: torch.Tensor,
+    cfg: PolicyConfig,
+) -> torch.Tensor:
+    """Apply temperature scaling to logits.
+
+    Args:
+        logits: Logits tensor to scale
+        cfg: Configuration dictionary containing generation settings
+
+    Returns:
+        torch.Tensor: Temperature-scaled logits
+    """
+    if "generation" in cfg and cfg["generation"] is not None:
+        logits.div_(cfg["generation"]["temperature"])
+    return logits
+
+
+def redistribute_logits_for_cp(
+    logits: torch.Tensor,
+    device_mesh: Any,
+    cp_mesh: Any,  # noqa: ARG001
+    sequence_dim: int = 1,
+) -> DTensor:
+    """Redistribute logits for context parallel processing.
+
+    Handles the case where logits may be TP-sharded DTensor or regular tensor,
+    and converts them to CP+TP sharded DTensor.
+
+    Args:
+        logits: Logits tensor (may be DTensor or regular tensor)
+        device_mesh: Full device mesh
+        cp_mesh: Context parallel mesh (kept for signature compatibility)
+        sequence_dim: Dimension for sequence sharding
+
+    Returns:
+        DTensor sharded on both CP and TP dimensions
+    """
+    if isinstance(logits, DTensor):
+        # Must be tp sharded
+        assert (
+            logits.device_mesh.ndim == 1
+            and logits.device_mesh.mesh_dim_names[0] == "tp"
+        ), "logits must be tp sharded"
+
+        # CP is implicitly sharded on the seq dim, so we need to redistribute to the tp dim
+        logits = DTensor.from_local(
+            logits.to_local(),
+            device_mesh=device_mesh[("cp", "tp")],
+            placements=[Shard(sequence_dim), Shard(-1)],
+        )
+    else:
+        logits = DTensor.from_local(
+            logits,
+            device_mesh=device_mesh[("cp", "tp")],
+            placements=[Shard(sequence_dim), Shard(-1)],
+        )
+    return logits
+
+
+def prepare_data_for_cp(
+    mb: BatchedDataDict[Any],
+    processed_inputs: ProcessedInputs,
+    cp_mesh: Any,
+    sequence_dim: int = 1,
+) -> tuple[torch.Tensor, BatchedDataDict[Any]]:
+    """Prepare data for context parallel processing.
+
+    Converts seq_index to full tensor and wraps CP-sharded tensors in DTensor.
+
+    Args:
+        mb: Microbatch data dictionary
+        processed_inputs: Processed inputs containing CP buffers
+        cp_mesh: Context parallel mesh
+        sequence_dim: Dimension for sequence sharding
+
+    Returns:
+        Tuple of (seq_index_dtensor, updated_mb)
+    """
+    seq_index_dtensor = (
+        DTensor.from_local(
+            processed_inputs.seq_index,
+            device_mesh=cp_mesh,
+            placements=[Shard(1)],
+        )
+        .full_tensor()
+        .squeeze(0)
+    )
+
+    mb["seq_index"] = seq_index_dtensor
+
+    for tensor_name in mb:
+        current_tensor = mb[tensor_name]
+        for buffer in processed_inputs.cp_buffers:
+            if current_tensor is buffer:
+                assert type(current_tensor) == torch.Tensor, (
+                    f"tensor {tensor_name} is not a tensor"
+                )
+                mb[tensor_name] = DTensor.from_local(
+                    current_tensor,
+                    device_mesh=cp_mesh,
+                    placements=[Shard(sequence_dim)],
+                )
+                break
+
+    return seq_index_dtensor, mb
+
+
+def forward_with_post_processing_fn(
+    model: nn.Module,
+    cfg: PolicyConfig,
+    post_processing_fn: PostProcessingFunction,
+    processed_mb: ProcessedMicrobatch,
+    is_reward_model: bool = False,
+    allow_flash_attn_args: bool = True,
+    global_valid_seqs: Optional[torch.Tensor] = None,
+    global_valid_toks: Optional[torch.Tensor] = None,
+    sequence_dim: int = 1,
+) -> Tuple[Any, dict[str, Any], ProcessedMicrobatch]:
+    """Perform forward pass with pre-processed microbatch and apply post-processing.
+
+    This function takes a pre-processed microbatch (with sequence packing already handled),
+    runs the forward step through the model, and applies the post-processing function
+    to compute the result.
+
+    Unlike the megatron approach which returns a callable, this directly computes
+    and returns the result since automodel uses PyTorch autograd.
+
+    Args:
+        model: The model to run forward pass on
+        cfg: Configuration dictionary
+        post_processing_fn: Post-processing function to apply to the logits
+        processed_mb: Pre-fetched ProcessedMicrobatch containing data and processed inputs
+        is_reward_model: Whether this is a reward model
+        allow_flash_attn_args: Whether to pass flash_attn_kwargs to model
+        global_valid_seqs: Global valid sequence count for loss normalization
+        global_valid_toks: Global valid token count for loss normalization
+        sequence_dim: Sequence dimension
+
+    Returns:
+        tuple: (result, metrics, processed_microbatch)
+            - result: Output from post-processing (loss, logprobs, topk, or scores)
+            - metrics: Dictionary of metrics from post-processing
+            - processed_microbatch: The ProcessedMicrobatch that was processed
+    """
+    # Extract the processed components
+    data_dict = processed_mb.data_dict
+    processed_inputs = processed_mb.processed_inputs
+
+    # Model forward pass
+    outputs = model_forward(
+        model,
+        processed_inputs,
+        is_reward_model=is_reward_model,
+        allow_flash_attn_args=allow_flash_attn_args,
+    )
+
+    # Extract logits from model outputs
+    logits = extract_logits(model, outputs)
+    del outputs
+
+    # Apply temperature scaling only for sampling-oriented post-processors
+    # Loss and score computations should use unscaled logits
+    if isinstance(post_processing_fn, (LogprobsPostProcessor, TopkLogitsPostProcessor)):
+        logits = apply_temperature_scaling(logits, cfg)
+
+    # Apply the post-processing function directly based on type
+    if isinstance(post_processing_fn, LossPostProcessor):
+        result, metrics = post_processing_fn(
+            logits=logits,
+            mb=data_dict,
+            processed_inputs=processed_inputs,
+            global_valid_seqs=global_valid_seqs,
+            global_valid_toks=global_valid_toks,
+            sequence_dim=sequence_dim,
+        )
+    elif isinstance(
+        post_processing_fn, (LogprobsPostProcessor, TopkLogitsPostProcessor)
+    ):
+        result = post_processing_fn(
+            logits=logits,
+            processed_inputs=processed_inputs,
+            input_lengths=data_dict["input_lengths"],
+            original_batch_size=processed_mb.original_batch_size,
+            original_seq_len=processed_mb.original_seq_len,
+            enable_seq_packing=post_processing_fn.enable_seq_packing,
+            sequence_dim=sequence_dim,
+        )
+        if isinstance(post_processing_fn, LogprobsPostProcessor):
+            metrics = {"logprobs": result}
+        else:
+            vals, idx = result
+            metrics = {"topk_logits": vals, "topk_indices": idx}
+    elif isinstance(post_processing_fn, ScorePostProcessor):
+        result = post_processing_fn(logits=logits)
+        metrics = {"scores": result}
+    else:
+        raise TypeError(
+            f"Unknown post-processing function type: {type(post_processing_fn)}"
+        )
+
+    del logits
+    return result, metrics, processed_mb
+
+
+def automodel_forward_backward(
+    model: nn.Module,
+    cfg: PolicyConfig,
+    data_iterator: Iterator[ProcessedMicrobatch],
+    post_processing_fn: PostProcessingFunction,
+    forward_only: bool = False,
+    is_reward_model: bool = False,
+    allow_flash_attn_args: bool = True,
+    global_valid_seqs: Optional[torch.Tensor] = None,
+    global_valid_toks: Optional[torch.Tensor] = None,
+    sequence_dim: int = 1,
+    dp_size: int = 1,
+    cp_size: int = 1,
+    num_global_batches: int = 1,
+    train_context_fn: Optional[Callable[[ProcessedInputs], Any]] = None,
+    num_valid_microbatches: Optional[int] = None,
+    on_microbatch_start: Optional[Callable[[int], None]] = None,
+) -> list[Tuple[Any, dict[str, Any]]]:
+    """Execute forward and backward passes for automodel.
+
+    This is the main training loop function that coordinates forward and backward
+    passes across multiple microbatches using PyTorch autograd.
+
+    Unlike megatron_forward_backward which uses Megatron's pipeline parallel
+    framework, this uses standard PyTorch operations.
+
+    Args:
+        model: The model to train
+        cfg: Configuration dictionary
+        data_iterator: Iterator yielding ProcessedMicrobatch objects (already processed)
+        num_microbatches: Number of microbatches to process
+        post_processing_fn: Post-processing function to apply to the logits
+        forward_only: If True, skip backward pass
+        is_reward_model: Whether this is a reward model
+        allow_flash_attn_args: Whether to pass flash_attn_kwargs to model
+        global_valid_seqs: Global valid sequence count for loss normalization
+        global_valid_toks: Global valid token count for loss normalization
+        sequence_dim: Sequence dimension
+        dp_size: Data parallel size
+        cp_size: Context parallel size
+        num_global_batches: Number of global batches (for metric scaling)
+        train_context_fn: Optional callable that takes ProcessedInputs and returns
+            a context manager for the forward/backward pass. If None, no context is used.
+        num_valid_microbatches: Number of valid (non-dummy) microbatches. If provided,
+            microbatches beyond this index are treated as dummy batches (loss *= 0).
+            If None, all microbatches are considered valid.
+        on_microbatch_start: Optional callback called at the start of each microbatch
+            with the microbatch index. Useful for cache clearing, etc.
+
+    Returns:
+        List of (result, metrics) tuples from each microbatch
+    """
+    from contextlib import nullcontext
+
+    results = []
+
+    for mb_idx, processed_mb in enumerate(data_iterator):
+        # Call optional callback at start of microbatch
+        if on_microbatch_start is not None:
+            on_microbatch_start(mb_idx)
+
+        processed_inputs = processed_mb.processed_inputs
+
+        # Create train context if factory provided, otherwise use nullcontext
+        if train_context_fn is not None:
+            ctx = train_context_fn(processed_inputs)
+        else:
+            ctx = nullcontext()
+
+        with ctx:
+            # Forward pass with post-processing
+            result, metrics, _ = forward_with_post_processing_fn(
+                model=model,
+                cfg=cfg,
+                post_processing_fn=post_processing_fn,
+                processed_mb=processed_mb,
+                is_reward_model=is_reward_model,
+                allow_flash_attn_args=allow_flash_attn_args,
+                global_valid_seqs=global_valid_seqs,
+                global_valid_toks=global_valid_toks,
+                sequence_dim=sequence_dim,
+            )
+
+            # Check if this is a dummy batch
+            is_dummy = (
+                num_valid_microbatches is not None and mb_idx >= num_valid_microbatches
+            )
+
+            # Scale metrics for aggregation (only for loss)
+            if isinstance(post_processing_fn, LossPostProcessor):
+                # skip the update for dummy batches
+                if not is_dummy:
+                    ## scale by the number of global batches so we get the correct
+                    ## value when summing metrics across all microbatches
+                    for k in metrics.keys():
+                        if "_min" in k or "_max" in k:
+                            continue
+
+                        metrics[k] /= num_global_batches
+                else:
+                    # Zero out loss for dummy batches
+                    result = result * 0
+
+                # Backward pass if training
+                if not forward_only:
+                    ## NOTE: invalid samples should be multiplied
+                    ## by zero in the loss function to prevent them
+                    ## from affecting the gradient calculation
+
+                    # when FSDP reduces the gradients over the DP dim, they're automatically averaged
+                    # but we want to sum them so we cancel out the average here
+                    loss = result * dp_size * cp_size
+                    loss.backward()
+
+        results.append((result, metrics))
+
+    return results
+
+
+class LossPostProcessor:
+    """Post-processor for computing training loss from model outputs."""
+
+    def __init__(
+        self,
+        loss_fn: LossFunction,
+        cfg: PolicyConfig,
+        device_mesh: Any,
+        cp_mesh: Any,
+        tp_mesh: Any,
+        cp_size: int,
+        dp_size: int,
+    ):
+        """Initialize LossPostProcessor.
+
+        Args:
+            loss_fn: Loss function to compute loss
+            cfg: Configuration dictionary
+            device_mesh: Full device mesh
+            cp_mesh: Context parallel mesh
+            tp_mesh: Tensor parallel mesh
+            cp_size: Context parallel size
+            dp_size: Data parallel size
+        """
+        self.loss_fn = loss_fn
+        self.cfg = cfg
+        self.device_mesh = device_mesh
+        self.cp_mesh = cp_mesh
+        self.tp_mesh = tp_mesh
+        self.cp_size = cp_size
+        self.dp_size = dp_size
+
+    def __call__(
+        self,
+        logits: torch.Tensor,
+        mb: BatchedDataDict[Any],
+        processed_inputs: ProcessedInputs,
+        global_valid_seqs: torch.Tensor,
+        global_valid_toks: torch.Tensor,
+        sequence_dim: int = 1,
+    ) -> tuple[torch.Tensor, dict[str, Any]]:
+        """Compute loss from logits.
+
+        Args:
+            logits: Model output logits
+            mb: Microbatch data
+            processed_inputs: Processed inputs
+            global_valid_seqs: Global valid sequence count
+            global_valid_toks: Global valid token count
+            sequence_dim: Sequence dimension
+
+        Returns:
+            Tuple of (loss, metrics)
+        """
+        # Handle CP redistribution
+        if self.cp_size > 1:
+            _, mb = prepare_data_for_cp(
+                mb, processed_inputs, self.cp_mesh, sequence_dim
+            )
+            logits = redistribute_logits_for_cp(
+                logits, self.device_mesh, self.cp_mesh, sequence_dim
+            )
+
+        # Wrap loss function for sequence packing if needed
+        if processed_inputs.has_flash_attention:
+            loss_fn_ = SequencePackingLossWrapper(
+                loss_fn=self.loss_fn,
+                cu_seqlens_q=processed_inputs.flash_attn_kwargs.cu_seqlens_q,
+                cu_seqlens_q_padded=processed_inputs.flash_attn_kwargs.cu_seqlens_q,
+            )
+        else:
+            loss_fn_ = self.loss_fn
+
+        loss, loss_metrics = loss_fn_(
+            logits,
+            mb,
+            global_valid_seqs,
+            global_valid_toks,
+        )
+
+        return loss, loss_metrics
+
+
+class LogprobsPostProcessor:
+    """Post-processor for computing log probabilities from model outputs."""
+
+    def __init__(
+        self,
+        cfg: PolicyConfig,
+        device_mesh: Any,
+        cp_mesh: Any,
+        tp_mesh: Any,
+        cp_size: int,
+        enable_seq_packing: bool = False,
+    ):
+        """Initialize LogprobsPostProcessor.
+
+        Args:
+            cfg: Configuration dictionary
+            device_mesh: Full device mesh
+            cp_mesh: Context parallel mesh
+            tp_mesh: Tensor parallel mesh
+            cp_size: Context parallel size
+            enable_seq_packing: Whether sequence packing is enabled
+        """
+        self.cfg = cfg
+        self.device_mesh = device_mesh
+        self.cp_mesh = cp_mesh
+        self.tp_mesh = tp_mesh
+        self.cp_size = cp_size
+        self.enable_seq_packing = enable_seq_packing
+        self.logprob_chunk_size = cfg.get("logprob_chunk_size", None)
+
+    def __call__(
+        self,
+        logits: torch.Tensor,
+        processed_inputs: ProcessedInputs,
+        input_lengths: torch.Tensor,
+        original_batch_size: int,
+        original_seq_len: int,
+        enable_seq_packing: bool,
+        sequence_dim: int = 1,
+    ) -> torch.Tensor:
+        """Compute token log probabilities from logits.
+
+        Args:
+            logits: Model output logits
+            processed_inputs: Processed inputs
+            input_lengths: Sequence lengths
+            original_batch_size: Original batch size before packing
+            original_seq_len: Original sequence length before packing
+            enable_seq_packing: Whether sequence packing is enabled
+            sequence_dim: Sequence dimension
+
+        Returns:
+            Token log probabilities tensor [batch_size, seq_length]
+        """
+        seq_len = processed_inputs.seq_len
+
+        if self.cp_size > 1:
+            seq_index_tensor = (
+                DTensor.from_local(
+                    processed_inputs.seq_index,
+                    device_mesh=self.cp_mesh,
+                    placements=[Shard(1)],
+                )
+                .full_tensor()
+                .squeeze(0)
+            )
+
+            input_ids_dtensor = DTensor.from_local(
+                processed_inputs.input_ids,
+                device_mesh=self.cp_mesh,
+                placements=[Shard(sequence_dim)],
+            )
+
+            logits = redistribute_logits_for_cp(
+                logits, self.device_mesh, self.cp_mesh, sequence_dim
+            )
+
+            token_logprobs = get_logprobs_from_vocab_parallel_logits(
+                logits,
+                input_ids_dtensor,
+                seq_index_tensor,
+                chunk_size=self.logprob_chunk_size,
+            )
+
+            assert token_logprobs.shape[1] == seq_len - 1
+        else:
+            if isinstance(logits, DTensor):
+                token_logprobs = get_logprobs_from_vocab_parallel_logits(
+                    logits,
+                    processed_inputs.input_ids,
+                    chunk_size=self.logprob_chunk_size,
+                )
+            else:
+                token_logprobs = self._compute_local_logprobs(
+                    logits, processed_inputs.input_ids
+                )
+
+        # Prepend 0 for first token to maintain sequence length
+        token_logprobs = torch.cat(
+            [torch.zeros_like(token_logprobs[:, :1]), token_logprobs], dim=1
+        )
+
+        # Handle sequence packing unpacking or mask application
+        if enable_seq_packing:
+            unpacked_logprobs = torch.zeros(
+                (original_batch_size, original_seq_len),
+                dtype=token_logprobs.dtype,
+                device=token_logprobs.device,
+            )
+            cu_seqlens = processed_inputs.flash_attn_kwargs.cu_seqlens_q
+            for i in range(original_batch_size):
+                start = cu_seqlens[i].item() + 1
+                end = cu_seqlens[i + 1].item()
+                seq_len_actual = input_lengths[i].item()
+                unpacked_logprobs[i, 1:seq_len_actual] = token_logprobs[0, start:end]
+            token_logprobs = unpacked_logprobs
+        else:
+            # Apply mask to zero out padding tokens logprobs
+            batch_size = processed_inputs.input_ids.shape[0]
+            post_attention_mask = torch.zeros(
+                (batch_size, seq_len),
+                dtype=torch.bool,
+                device=token_logprobs.device,
+            )
+            for i, length in enumerate(input_lengths):
+                # For right-padded sequence, set 1s at the beginning of the sequence
+                post_attention_mask[i, :length] = 1
+            token_logprobs = token_logprobs * post_attention_mask
+
+        return token_logprobs
+
+    def _compute_local_logprobs(
+        self,
+        logits: torch.Tensor,
+        input_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute logprobs locally without distributed processing.
+
+        Args:
+            logits: Model output logits
+            input_ids: Input token IDs
+
+        Returns:
+            Token log probabilities
+        """
+        if self.logprob_chunk_size is not None:
+            logits_seq_len = int(logits.shape[1])
+            num_chunks = (
+                logits_seq_len + self.logprob_chunk_size - 1
+            ) // self.logprob_chunk_size
+            chunked_log_probs = []
+            for chunk_idx in range(num_chunks):
+                chunk_start = chunk_idx * self.logprob_chunk_size
+                chunk_end = min(
+                    logits_seq_len,
+                    (chunk_idx + 1) * self.logprob_chunk_size,
+                )
+                chunk_logits = logits[:, chunk_start:chunk_end, :].to(torch.float32)
+                log_probs = torch.nn.functional.log_softmax(chunk_logits, dim=-1)
+                chunked_log_probs.append(log_probs)
+            log_probs = torch.cat(chunked_log_probs, dim=1)
+            del chunked_log_probs
+        else:
+            logits = logits.to(torch.float32)
+            log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
+
+        # Extract logprobs for each token in the sequence by gathering the logprob
+        # corresponding to the next token at each position
+        # Input shapes:
+        #   log_probs: [batch_size, sequence_length, vocab_size] - logits for each position
+        #   token_ids: [batch_size, sequence_length] - actual tokens
+        # Output shape: [batch_size, sequence_length] - logprob of each token given previous
+        # We get logprob of token[t+1] from logits[t], prepending 0 to maintain sequence length
+        next_tokens = input_ids[:, 1:]
+        log_probs = log_probs[:, :-1]
+        token_logprobs = log_probs.gather(
+            dim=-1, index=next_tokens.unsqueeze(-1)
+        ).squeeze(-1)
+        del log_probs
+
+        return token_logprobs
+
+
+class TopkLogitsPostProcessor:
+    """Post-processor for computing top-k logits from model outputs."""
+
+    def __init__(
+        self,
+        cfg: PolicyConfig,
+        device_mesh: Any,
+        cp_mesh: Any,
+        tp_mesh: Any,
+        cp_size: int,
+        k: int,
+        enable_seq_packing: bool = False,
+    ):
+        """Initialize TopkLogitsPostProcessor.
+
+        Args:
+            cfg: Configuration dictionary
+            device_mesh: Full device mesh
+            cp_mesh: Context parallel mesh
+            tp_mesh: Tensor parallel mesh
+            cp_size: Context parallel size
+            k: Number of top logits to return
+            enable_seq_packing: Whether sequence packing is enabled
+        """
+        self.cfg = cfg
+        self.device_mesh = device_mesh
+        self.cp_mesh = cp_mesh
+        self.tp_mesh = tp_mesh
+        self.cp_size = cp_size
+        self.k = k
+        self.enable_seq_packing = enable_seq_packing
+
+    def __call__(
+        self,
+        logits: torch.Tensor,
+        processed_inputs: ProcessedInputs,
+        input_lengths: torch.Tensor,
+        original_batch_size: int,
+        original_seq_len: int,
+        enable_seq_packing: bool,
+        sequence_dim: int = 1,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compute top-k logits and indices from model outputs.
+
+        Args:
+            logits: Model output logits
+            processed_inputs: Processed inputs
+            input_lengths: Sequence lengths
+            original_batch_size: Original batch size before packing
+            original_seq_len: Original sequence length before packing
+            enable_seq_packing: Whether sequence packing is enabled
+            sequence_dim: Sequence dimension
+
+        Returns:
+            Tuple of (top-k values, top-k indices) tensors
+        """
+        if self.cp_size > 1:
+            logits = redistribute_logits_for_cp(
+                logits, self.device_mesh, self.cp_mesh, sequence_dim
+            )
+
+            # Deal with TP first
+            local_logits = logits.to_local()  # [B, S_cp, V_tp]
+
+            tp_group = self.tp_mesh.get_group()
+            tp_rank = torch.distributed.get_rank(tp_group)
+            V_local = int(local_logits.shape[-1])
+            vocab_start_index = tp_rank * V_local
+            vocab_end_index = (tp_rank + 1) * V_local
+
+            vals, idx = distributed_vocab_topk(
+                local_logits,
+                k=self.k,
+                tp_group=tp_group,
+                vocab_start_index=vocab_start_index,
+                vocab_end_index=vocab_end_index,
+            )
+            # [B, S_cp, k]
+
+            cp_group = self.cp_mesh.get_group()
+
+            vals = allgather_cp_sharded_tensor(vals, cp_group, seq_dim=sequence_dim)
+            idx = allgather_cp_sharded_tensor(idx, cp_group, seq_dim=sequence_dim)
+            # [B, S, k]
+        else:
+            # Compute top-k over full sequence length
+            if isinstance(logits, DTensor):
+                local_logits = logits.to_local()  # [B, S, V_local]
+                tp_group = self.tp_mesh.get_group()
+                tp_rank = torch.distributed.get_rank(tp_group)
+                V_local = int(local_logits.shape[-1])
+                vocab_start_index = tp_rank * V_local
+                vocab_end_index = (tp_rank + 1) * V_local
+
+                vals, idx = distributed_vocab_topk(
+                    local_logits,
+                    k=self.k,
+                    tp_group=tp_group,
+                    vocab_start_index=vocab_start_index,
+                    vocab_end_index=vocab_end_index,
+                )
+            else:
+                full_logits = logits.to(torch.float32)
+                vals, idx = torch.topk(full_logits, k=self.k, dim=-1)
+
+        # Handle sequence packing unpacking
+        if enable_seq_packing:
+            # Unpack top-k results from packed format back to original batch format
+            # vals: [1, packed_seq_len, k] -> [original_batch_size, original_seq_len, k]
+            # idx: [1, packed_seq_len, k] -> [original_batch_size, original_seq_len, k]
+            unpacked_vals = torch.zeros(
+                (original_batch_size, original_seq_len, self.k),
+                dtype=vals.dtype,
+                device=vals.device,
+            )
+            unpacked_idx = torch.zeros(
+                (original_batch_size, original_seq_len, self.k),
+                dtype=idx.dtype,
+                device=idx.device,
+            )
+
+            cu_seqlens = processed_inputs.flash_attn_kwargs.cu_seqlens_q
+
+            for i in range(original_batch_size):
+                start = cu_seqlens[i].item()
+                end = cu_seqlens[i + 1].item()
+                seq_len_actual = input_lengths[i].item()
+
+                # Extract the corresponding portion from packed results
+                # Note: vals and idx are [1, packed_seq_len, k] due to packing
+                unpacked_vals[i, :seq_len_actual, :] = vals[0, start:end, :]
+                unpacked_idx[i, :seq_len_actual, :] = idx[0, start:end, :]
+
+            vals = unpacked_vals
+            idx = unpacked_idx
+
+        return vals, idx
+
+
+class ScorePostProcessor:
+    """Post-processor for computing reward model scores from model outputs."""
+
+    def __init__(
+        self,
+        cfg: PolicyConfig,
+    ):
+        """Initialize ScorePostProcessor.
+
+        Args:
+            cfg: Configuration dictionary
+        """
+        self.cfg = cfg
+
+    def __call__(
+        self,
+        logits: torch.Tensor,
+    ) -> torch.Tensor:
+        """Extract scores from reward model outputs.
+
+        Args:
+            logits: Model output logits
+
+        Returns:
+            Scores tensor
+        """
+        logits = logits.to(torch.float32)
+        rm_scores = to_local_if_dtensor(logits)
+        rm_scores = rm_scores.squeeze(-1)
+
+        return rm_scores
+
+
+def aggregate_training_statistics(
+    losses: list[float],
+    all_mb_metrics: list[dict[str, Any]],
+    grad_norm: Optional[torch.Tensor],
+    dp_group: Any,
+    dtype: torch.dtype,
+) -> dict[str, Any]:
+    """Aggregate training statistics across microbatches and ranks.
+
+    Args:
+        losses: List of loss values from each microbatch
+        all_mb_metrics: List of metrics dictionaries from each microbatch
+        grad_norm: Gradient norm tensor (or None if eval mode)
+        dp_group: Data parallel process group for all-reduce
+        dtype: Model dtype for metrics
+
+    Returns:
+        Dictionary containing aggregated metrics including global_loss, grad_norm, etc.
+    """
+    # Compute global loss across all ranks
+    with torch.no_grad():
+        global_loss = torch.tensor(losses, device="cuda")
+        torch.distributed.all_reduce(global_loss, group=dp_group)
+
+    # Aggregate metrics across all microbatches
+    mb_metrics = defaultdict(list)
+    for m in all_mb_metrics:
+        for k, v in m.items():
+            mb_metrics[k].append(v)
+
+    metrics = {
+        "global_loss": global_loss.cpu(),
+        "grad_norm": grad_norm,
+        "rank": torch.distributed.get_rank(),
+        "gpu_name": torch.cuda.get_device_name(),
+        "model_dtype": dtype,
+        "all_mb_metrics": dict(mb_metrics),
+    }
+
+    return metrics
diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py b/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py
index 9342c5c138..94d27277c1 100644
--- a/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py
+++ b/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py
@@ -15,7 +15,6 @@
 import contextlib
 import gc
 import warnings
-from collections import defaultdict
 from contextlib import AbstractContextManager, contextmanager, nullcontext
 from typing import Any, Generator, Optional
 
@@ -34,21 +33,16 @@
 )
 from nemo_automodel.components.training.utils import scale_grads_and_clip_grad_norm
 from torch import nn
-from torch.distributed.tensor import DTensor, Shard
+from torch.distributed.tensor import DTensor
 from transformers import (
     AutoProcessor,
     AutoTokenizer,
 )
 
 from nemo_rl.algorithms.interfaces import LossFunction
-from nemo_rl.algorithms.loss_functions import SequencePackingLossWrapper
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
-from nemo_rl.distributed.model_utils import (
-    allgather_cp_sharded_tensor,
-    distributed_vocab_topk,
-    get_logprobs_from_vocab_parallel_logits,
-)
 from nemo_rl.models.automodel.data import (
+    check_sequence_dim,
     get_microbatch_iterator,
     process_global_batch,
 )
@@ -58,6 +52,15 @@
     setup_reference_model_state,
     validate_and_prepare_config,
 )
+from nemo_rl.models.automodel.train import (
+    LogprobsPostProcessor,
+    LossPostProcessor,
+    ScorePostProcessor,
+    TopkLogitsPostProcessor,
+    aggregate_training_statistics,
+    automodel_forward_backward,
+    forward_with_post_processing_fn,
+)
 from nemo_rl.models.policy import PolicyConfig
 from nemo_rl.models.policy.interfaces import (
     ColocatablePolicyInterface,
@@ -320,11 +323,6 @@ def __init__(
             _runtime_is_reward_model,  # Duplicate, already set as _is_reward_model
         ) = runtime_config
 
-    def _apply_temperature_scaling(self, logits: torch.Tensor) -> torch.Tensor:
-        if "generation" in self.cfg and self.cfg["generation"] is not None:
-            logits.div_(self.cfg["generation"]["temperature"])
-        return logits
-
     @wrap_with_nvtx_name("dtensor_policy_worker_v2/train")
     def train(
         self,
@@ -348,14 +346,8 @@ def train(
         )
         num_global_batches = int(total_dataset_size.item()) // gbs
 
-        # dim 1 is always assumed to be the sequence dim, sanity check this here
-        sequence_dim = 1
-        seq_dim_size = data.get("input_ids").shape[sequence_dim]
-        for k, v in data.items():
-            if torch.is_tensor(v) and len(v.shape) > 1:
-                assert v.shape[sequence_dim] == seq_dim_size, (
-                    f"Dim 1 must be the sequence dim, expected dim 1={seq_dim_size} but got shape {v.shape}"
-                )
+        # Validate sequence dimension
+        sequence_dim, _ = check_sequence_dim(data)
 
         if eval_mode:
             ctx: AbstractContextManager[Any] = torch.no_grad()
@@ -365,9 +357,44 @@ def train(
             # Ensure model is in training mode
             self.model.train()
 
+        # Create loss post-processor
+        loss_post_processor = LossPostProcessor(
+            loss_fn=loss_fn,
+            cfg=self.cfg,
+            device_mesh=self.device_mesh,
+            cp_mesh=self.cp_mesh,
+            tp_mesh=self.tp_mesh,
+            cp_size=self.cp_size,
+            dp_size=self.dp_size,
+        )
+
+        # Create train context factory
+        def train_context_fn(processed_inputs):
+            return get_train_context(
+                cp_size=self.cp_size,
+                cp_mesh=self.cp_mesh,
+                cp_buffers=processed_inputs.cp_buffers,
+                sequence_dim=sequence_dim,
+                dtype=self.dtype,
+                autocast_enabled=self.autocast_enabled,
+            )
+
+        # Setup cache clearing callback if configured
+        empty_cache_steps = self.cfg.get("dtensor_cfg", {}).get(
+            "clear_cache_every_n_steps"
+        )
+        if empty_cache_steps:
+            warnings.warn(
+                f"Emptying cache every {empty_cache_steps} microbatches; doing so unnecessarily would incur a large performance overhead.",
+            )
+
+        def on_microbatch_start(mb_idx):
+            if empty_cache_steps and mb_idx % empty_cache_steps == 0:
+                torch.cuda.empty_cache()
+
         with ctx:
             # Get data from batch and move to device
-            data.to("cuda")
+            data = data.to("cuda")
 
             losses = []
             all_mb_metrics = []
@@ -385,7 +412,7 @@ def train(
                 global_valid_toks = gb_result["global_valid_toks"]
 
                 self.optimizer.zero_grad()
-                mb_losses = []
+
                 # Get microbatch iterator based on batching strategy
                 processed_iterator, iterator_len = get_microbatch_iterator(
                     batch,
@@ -396,175 +423,39 @@ def train(
                     cp_size=self.cp_size,
                 )
 
-                empty_cache_steps = self.cfg.get("dtensor_cfg", {}).get(
-                    "clear_cache_every_n_steps"
+                # Use automodel_forward_backward for the training loop
+                mb_results = automodel_forward_backward(
+                    model=self.model,
+                    cfg=self.cfg,
+                    data_iterator=processed_iterator,
+                    post_processing_fn=loss_post_processor,
+                    forward_only=eval_mode,
+                    is_reward_model=self._is_reward_model,
+                    allow_flash_attn_args=self.allow_flash_attn_args,
+                    global_valid_seqs=global_valid_seqs,
+                    global_valid_toks=global_valid_toks,
+                    sequence_dim=sequence_dim,
+                    dp_size=self.dp_size,
+                    cp_size=self.cp_size,
+                    num_global_batches=num_global_batches,
+                    train_context_fn=train_context_fn,
+                    num_valid_microbatches=iterator_len,
+                    on_microbatch_start=on_microbatch_start,
                 )
-                if empty_cache_steps:
-                    warnings.warn(
-                        f"Emptying cache every {empty_cache_steps} microbatches, doing so unnnecessarily would incur a large performance overhead."
-                    )
 
-                for mb_idx, processed_mb in enumerate(processed_iterator):
-                    # Conditioanlly empty cache when sensitive to fragmentation
-                    if empty_cache_steps and mb_idx % empty_cache_steps == 0:
-                        torch.cuda.empty_cache()
-
-                    # Extract data dict and processed inputs
-                    mb = processed_mb.data_dict
-                    processed_inputs = processed_mb.processed_inputs
-
-                    # Extract values from processed inputs for use in forward pass
-                    input_ids = processed_inputs.input_ids
-                    attention_mask = processed_inputs.attention_mask
-                    position_ids = processed_inputs.position_ids
-                    flash_attn_kwargs = processed_inputs.flash_attn_kwargs
-                    vlm_kwargs = processed_inputs.vlm_kwargs
-                    cp_buffers = processed_inputs.cp_buffers
-                    seq_index = processed_inputs.seq_index
-                    seq_len = processed_inputs.seq_len
-
-                    # get_train_context handles both context parallel and autocast
-                    with get_train_context(
-                        cp_size=self.cp_size,
-                        cp_mesh=self.cp_mesh,
-                        cp_buffers=cp_buffers,
-                        sequence_dim=sequence_dim,
-                        dtype=self.dtype,
-                        autocast_enabled=self.autocast_enabled,
-                    ):
-                        model_args = dict(
-                            input_ids=input_ids,
-                            attention_mask=attention_mask,
-                            position_ids=position_ids,
-                            use_cache=False,
-                            flash_attn_kwargs=flash_attn_kwargs,
-                            **vlm_kwargs,
-                        )
-
-                        if self._is_reward_model:
-                            # `flash_attn_kwarg` is not supported for `LlamaForSequenceClassification`.
-                            # Note that it should be empty anyway since sequence packing
-                            # is not supported for reward models.
-                            assert not processed_inputs.has_flash_attention
-                            del model_args["flash_attn_kwargs"]
-                        # remove flash_attn_kwargs if there are multimodal kwargs
-                        if processed_inputs.is_multimodal:
-                            del model_args["flash_attn_kwargs"]
-
-                        if (
-                            not self.allow_flash_attn_args
-                            and "flash_attn_kwargs" in model_args
-                        ):
-                            del model_args["flash_attn_kwargs"]
-
-                        outputs = self.model(**model_args)
-
-                        # Get logprobs
-                        if isinstance(outputs, (torch.Tensor, DTensor)):
-                            # custom models (e.g., those coming from AutoModel) can output logits directly
-                            logits = outputs
-                        elif not hasattr(outputs, "logits"):
-                            logits = self.model.lm_head(outputs.last_hidden_state)
-                        else:
-                            logits = outputs.logits
-                        del outputs
-
-                        # Apply temperature scaling
-                        logits = self._apply_temperature_scaling(logits)
-
-                        if self.cp_size > 1:
-                            seq_index_dtensor = (
-                                DTensor.from_local(
-                                    seq_index,
-                                    device_mesh=self.cp_mesh,
-                                    placements=[Shard(1)],
-                                )
-                                .full_tensor()
-                                .squeeze(0)
-                            )
-
-                            mb["seq_index"] = seq_index_dtensor
-
-                            for tensor_name in mb:
-                                current_tensor = mb[tensor_name]
-                                for buffer in cp_buffers:
-                                    if current_tensor is buffer:
-                                        assert type(current_tensor) == torch.Tensor, (
-                                            f"tensor {tensor_name} is not a tensor"
-                                        )
-                                        mb[tensor_name] = DTensor.from_local(
-                                            current_tensor,
-                                            device_mesh=self.cp_mesh,
-                                            placements=[Shard(sequence_dim)],
-                                        )
-                                        break
-
-                            if isinstance(logits, DTensor):
-                                # Must be tp sharded
-                                assert (
-                                    logits.device_mesh.ndim == 1
-                                    and logits.device_mesh.mesh_dim_names[0] == "tp"
-                                ), "logits must be tp sharded"
-
-                                # CP is implicitly sharded on the seq dim, so we need to redistribute to the tp dim
-                                logits = DTensor.from_local(
-                                    logits.to_local(),
-                                    device_mesh=self.device_mesh[("cp", "tp")],
-                                    placements=[Shard(sequence_dim), Shard(-1)],
-                                )
-                            else:
-                                logits = DTensor.from_local(
-                                    logits,
-                                    device_mesh=self.device_mesh[("cp", "tp")],
-                                    placements=[Shard(sequence_dim), Shard(-1)],
-                                )
-
-                        if self.enable_seq_packing:
-                            loss_fn_ = SequencePackingLossWrapper(
-                                loss_fn=loss_fn,
-                                cu_seqlens_q=flash_attn_kwargs.cu_seqlens_q,
-                                cu_seqlens_q_padded=flash_attn_kwargs.cu_seqlens_q,
-                            )
-                        else:
-                            loss_fn_ = loss_fn
-
-                        loss, loss_metrics = loss_fn_(
-                            logits,
-                            mb,
-                            global_valid_seqs,
-                            global_valid_toks,
-                        )
-                        del logits
-
-                        # skip the update for dummy batches
-                        if mb_idx < iterator_len:
-                            ## scale by the number of global batches so we get the correct
-                            ## value when summing metrics across all microbatches
-                            for k in loss_metrics.keys():
-                                if "_min" in k or "_max" in k:
-                                    continue
-                                loss_metrics[k] /= num_global_batches
-                            num_valid_samples = loss_metrics["num_valid_samples"]
-                            loss_metrics["lr"] = self.optimizer.param_groups[0]["lr"]
-                            loss_metrics["global_valid_seqs"] = global_valid_seqs.item()
-                            loss_metrics["global_valid_toks"] = global_valid_toks.item()
-                        else:
-                            loss *= 0
-
-                        # Backward pass
-                        if not eval_mode:
-                            ## NOTE: invalid samples should be multiplied
-                            ## by zero in the loss function to prevent them
-                            ## from affecting the gradient calculation
-
-                            # when FSDP reduces the gradients over the DP dim, they're automatically averaged
-                            # but we want to sum them so we cancel out the average here
-                            loss *= self.dp_size * self.cp_size
-                            loss.backward()
-
-                    if num_valid_samples > 0:
-                        mb_losses.append(loss.item())
-                        all_mb_metrics.append(loss_metrics)
+                # Extract losses and metrics from results
+                mb_losses = []
+                for mb_idx, (loss, loss_metrics) in enumerate(mb_results):
+                    # Only process valid (non-dummy) batches for metrics
+                    if mb_idx < iterator_len:
+                        num_valid_samples = loss_metrics["num_valid_samples"]
+                        loss_metrics["lr"] = self.optimizer.param_groups[0]["lr"]
+                        loss_metrics["global_valid_seqs"] = global_valid_seqs.item()
+                        loss_metrics["global_valid_toks"] = global_valid_toks.item()
+
+                        if num_valid_samples > 0:
+                            mb_losses.append(loss.item())
+                            all_mb_metrics.append(loss_metrics)
 
                 grad_norm: Optional[float | torch.Tensor] = None
                 if not eval_mode:
@@ -602,30 +493,17 @@ def train(
             # the memory allocator before moving on
             torch.cuda.empty_cache()
 
-            # Compute global loss across all ranks
-            with torch.no_grad():
-                global_loss = torch.tensor(losses, device="cuda")
-                torch.distributed.all_reduce(
-                    global_loss, group=self.dp_mesh.get_group()
-                )
-            # Aggregate metrics across all microbatches
-            mb_metrics = defaultdict(list)
-            for m in all_mb_metrics:
-                for k, v in m.items():
-                    mb_metrics[k].append(v)
-
-            metrics = {
-                "global_loss": global_loss.cpu(),
-                "grad_norm": grad_norm,
-                "rank": torch.distributed.get_rank(),
-                "gpu_name": torch.cuda.get_device_name(),
-                "model_dtype": self.dtype,
-                "all_mb_metrics": dict(mb_metrics),
-            }
+            # Aggregate training statistics across microbatches and ranks
+            metrics = aggregate_training_statistics(
+                losses=losses,
+                all_mb_metrics=all_mb_metrics,
+                grad_norm=grad_norm,
+                dp_group=self.dp_mesh.get_group(),
+                dtype=self.dtype,
+            )
 
             return metrics
 
-    # TODO @Rayen Tian: Related Issue: Refactor shared logic between score() and get_logprobs() (https://github.com/NVIDIA-NeMo/RL/issues/1094)
     @wrap_with_nvtx_name("dtensor_policy_worker_v2/get_logprobs")
     def get_logprobs(
         self, data: BatchedDataDict[Any], micro_batch_size: Optional[int] = None
@@ -647,20 +525,23 @@ def get_logprobs(
             if micro_batch_size is not None
             else self.cfg["logprob_batch_size"]
         )
-        logprob_chunk_size = self.cfg.get("logprob_chunk_size", None)
-
-        # dim 1 is always assumed to be the sequence dim, sanity check this here
-        sequence_dim = 1
-        seq_dim_size = data.get("input_ids").shape[sequence_dim]
-        for k, v in data.items():
-            if torch.is_tensor(v) and len(v.shape) > 1:
-                assert v.shape[sequence_dim] == seq_dim_size, (
-                    f"Dim 1 must be the sequence dim, expected dim 1={seq_dim_size} but got shape {v.shape}"
-                )
+
+        # Validate sequence dimension
+        sequence_dim, seq_dim_size = check_sequence_dim(data)
 
         all_log_probs = []
         self.model.eval()
 
+        # Create logprobs post-processor
+        logprobs_post_processor = LogprobsPostProcessor(
+            cfg=self.cfg,
+            device_mesh=self.device_mesh,
+            cp_mesh=self.cp_mesh,
+            tp_mesh=self.tp_mesh,
+            cp_size=self.cp_size,
+            enable_seq_packing=self.enable_seq_packing,
+        )
+
         with torch.no_grad():
             data.to("cuda")
             # Get microbatch iterator based on batching strategy
@@ -673,195 +554,32 @@ def get_logprobs(
                 cp_size=self.cp_size,
             )
 
-            step = 0
             for batch_idx, processed_mb in enumerate(processed_iterator):
-                step += 1
-                # Extract data dict and processed inputs
-                lp_batch = processed_mb.data_dict
                 processed_inputs = processed_mb.processed_inputs
 
-                # Use original shapes from ProcessedMicrobatch (needed for unpacking later)
-                original_batch_size = processed_mb.original_batch_size
-                original_seq_len = processed_mb.original_seq_len
-
-                # Extract values from processed inputs
-                input_ids = processed_inputs.input_ids
-                attention_mask = processed_inputs.attention_mask
-                position_ids = processed_inputs.position_ids
-                flash_attn_kwargs = processed_inputs.flash_attn_kwargs
-                vlm_kwargs = processed_inputs.vlm_kwargs
-                cp_buffers = processed_inputs.cp_buffers
-                seq_index = processed_inputs.seq_index
-                seq_len = processed_inputs.seq_len
-
-                input_lengths = lp_batch.get("input_lengths")
-                batch_size = input_ids.shape[0]
-
-                # Create post_attention_mask for right-padded data (used for masking after forward)
-                if not self.enable_seq_packing:
-                    post_attention_mask = torch.zeros(
-                        (batch_size, seq_len), dtype=torch.bool, device=input_ids.device
-                    )
-                    for i, length in enumerate(input_lengths):
-                        # For right-padded sequence, set 1s at the beginning of the sequence
-                        post_attention_mask[i, :length] = 1
-
                 with get_train_context(
                     cp_size=self.cp_size,
                     cp_mesh=self.cp_mesh,
-                    cp_buffers=cp_buffers,
+                    cp_buffers=processed_inputs.cp_buffers,
                     sequence_dim=sequence_dim,
                     dtype=self.dtype,
                     autocast_enabled=self.autocast_enabled,
                 ):
-                    model_args = dict(
-                        input_ids=input_ids,
-                        attention_mask=attention_mask,
-                        position_ids=position_ids,
-                        use_cache=False,
-                        flash_attn_kwargs=flash_attn_kwargs,
-                        **vlm_kwargs,
+                    # Use forward_with_post_processing_fn for forward pass and post-processing
+                    token_logprobs, _metrics, _ = forward_with_post_processing_fn(
+                        model=self.model,
+                        cfg=self.cfg,
+                        post_processing_fn=logprobs_post_processor,
+                        processed_mb=processed_mb,
+                        is_reward_model=False,
+                        allow_flash_attn_args=self.allow_flash_attn_args,
+                        sequence_dim=sequence_dim,
                     )
-                    if processed_inputs.is_multimodal:
-                        del model_args["flash_attn_kwargs"]
-
-                    if (
-                        not self.allow_flash_attn_args
-                        and "flash_attn_kwargs" in model_args
-                    ):
-                        del model_args["flash_attn_kwargs"]
-
-                    outputs = self.model(**model_args)
-
-                    logits = outputs.logits if hasattr(outputs, "logits") else outputs
-
-                    # Apply temperature scaling
-                    logits = self._apply_temperature_scaling(logits)
-
-                    if self.cp_size > 1:
-                        seq_index_tensor = (
-                            DTensor.from_local(
-                                seq_index,
-                                device_mesh=self.cp_mesh,
-                                placements=[Shard(1)],
-                            )
-                            .full_tensor()
-                            .squeeze(0)
-                        )
-
-                        input_ids_dtensor = DTensor.from_local(
-                            input_ids,
-                            device_mesh=self.cp_mesh,
-                            placements=[Shard(sequence_dim)],
-                        )
-
-                        if isinstance(logits, DTensor):
-                            # Must be tp sharded
-                            assert (
-                                logits.device_mesh.ndim == 1
-                                and logits.device_mesh.mesh_dim_names[0] == "tp"
-                            ), "logits must be tp sharded"
-
-                            # CP is implicitly sharded on the seq dim, so we need to redistribute to the tp dim
-                            logits = DTensor.from_local(
-                                logits.to_local(),
-                                device_mesh=self.device_mesh[("cp", "tp")],
-                                placements=[Shard(sequence_dim), Shard(-1)],
-                            )
-                        else:
-                            logits = DTensor.from_local(
-                                logits,
-                                device_mesh=self.device_mesh[("cp", "tp")],
-                                placements=[Shard(sequence_dim), Shard(-1)],
-                            )
-
-                        token_logprobs = get_logprobs_from_vocab_parallel_logits(
-                            logits,
-                            input_ids_dtensor,
-                            seq_index_tensor,
-                            chunk_size=logprob_chunk_size,
-                        )
-
-                        assert token_logprobs.shape[1] == seq_len - 1
-                    else:
-                        if isinstance(logits, DTensor):
-                            token_logprobs = get_logprobs_from_vocab_parallel_logits(
-                                logits,
-                                input_ids,
-                                chunk_size=logprob_chunk_size,
-                            )
-                        else:
-                            if logprob_chunk_size is not None:
-                                logits_seq_len = int(logits.shape[1])
-                                num_chunks = (
-                                    logits_seq_len + logprob_chunk_size - 1
-                                ) // logprob_chunk_size
-                                chunked_log_probs = []
-                                for chunk_idx in range(num_chunks):
-                                    chunk_start = chunk_idx * logprob_chunk_size
-                                    chunk_end = min(
-                                        logits_seq_len,
-                                        (chunk_idx + 1) * logprob_chunk_size,
-                                    )
-                                    chunk_logits = logits[
-                                        :, chunk_start:chunk_end, :
-                                    ].to(torch.float32)
-                                    log_probs = torch.nn.functional.log_softmax(
-                                        chunk_logits, dim=-1
-                                    )
-                                    chunked_log_probs.append(log_probs)
-                                log_probs = torch.cat(chunked_log_probs, dim=1)
-                                del chunked_log_probs
-                            else:
-                                logits = logits.to(torch.float32)
-                                log_probs = torch.nn.functional.log_softmax(
-                                    logits, dim=-1
-                                )
-                            # Extract logprobs for each token in the sequence by gathering the logprob
-                            # corresponding to the next token at each position
-                            # Input shapes:
-                            #   log_probs: [batch_size, sequence_length, vocab_size] - logits for each position
-                            #   token_ids: [batch_size, sequence_length] - actual tokens
-                            # Output shape: [batch_size, sequence_length] - logprob of each token given previous
-                            # We get logprob of token[t+1] from logits[t], prepending 0 to maintain sequence length
-                            next_tokens = input_ids[:, 1:]
-                            log_probs = log_probs[:, :-1]
-                            token_logprobs = log_probs.gather(
-                                dim=-1, index=next_tokens.unsqueeze(-1)
-                            ).squeeze(-1)
-                            del log_probs
-
-                del outputs, logits
-
-                token_logprobs = torch.cat(
-                    [torch.zeros_like(token_logprobs[:, :1]), token_logprobs], dim=1
-                )
 
                 # skip keeping the logprobs for the dummy batches
                 if batch_idx >= iterator_len:
                     continue
 
-                if not self.enable_seq_packing:
-                    # Apply mask to zero out padding tokens logprobs
-                    token_logprobs = token_logprobs * post_attention_mask
-                else:
-                    # For packed sequences, unpack logprobs
-                    # Use original_batch_size since packed sequences have shape [1, packed_seq_len]
-                    unpacked_logprobs = torch.zeros(
-                        (original_batch_size, seq_dim_size),
-                        dtype=token_logprobs.dtype,
-                        device=token_logprobs.device,
-                    )
-                    cu_seqlens = flash_attn_kwargs.cu_seqlens_q
-                    for i in range(original_batch_size):
-                        start = cu_seqlens[i].item() + 1
-                        end = cu_seqlens[i + 1].item()
-                        seq_len_actual = input_lengths[i].item()
-                        unpacked_logprobs[i, 1:seq_len_actual] = token_logprobs[
-                            0, start:end
-                        ]
-                    token_logprobs = unpacked_logprobs
-
                 all_log_probs.append(token_logprobs)
 
         # Concatenate all batches
@@ -879,20 +597,19 @@ def get_logprobs(
 
         return return_data
 
-    # TODO @Rayen Tian: Related Issue: Refactor shared logic between score() and get_logprobs() (https://github.com/NVIDIA-NeMo/RL/issues/1094)
     @wrap_with_nvtx_name("dtensor_policy_worker_v2/score")
     def score(self, data: BatchedDataDict) -> BatchedDataDict[ScoreOutputSpec]:
         global_batch_size = min(self.cfg["batch_size"], data.size)
 
-        sequence_dim = 1
-        seq_dim_size = data.get("input_ids").shape[sequence_dim]
-        for k, v in data.items():
-            if torch.is_tensor(v) and len(v.shape) > 1:
-                assert v.shape[sequence_dim] == seq_dim_size, (
-                    f"Dim 1 must be the sequence dim, expected dim 1={seq_dim_size} but got shape {v.shape}"
-                )
+        # Validate sequence dimension
+        sequence_dim, _ = check_sequence_dim(data)
+
         self.model.eval()
         print("Begin to batch datas")
+
+        # Create score post-processor
+        score_post_processor = ScorePostProcessor(cfg=self.cfg)
+
         with torch.no_grad():
             data.to("cuda")
             # Get microbatch iterator based on batching strategy
@@ -905,52 +622,28 @@ def score(self, data: BatchedDataDict) -> BatchedDataDict[ScoreOutputSpec]:
                 cp_size=self.cp_size,
             )
 
-            step = 0
             all_rm_scores = []
             for batch_idx, processed_mb in enumerate(processed_iterator):
-                step += 1
-                # Extract processed inputs
                 processed_inputs = processed_mb.processed_inputs
 
-                # Extract values from processed inputs
-                input_ids = processed_inputs.input_ids
-                attention_mask = processed_inputs.attention_mask
-                position_ids = processed_inputs.position_ids
-                flash_attn_kwargs = processed_inputs.flash_attn_kwargs
-                vlm_kwargs = processed_inputs.vlm_kwargs
-                cp_buffers = processed_inputs.cp_buffers
-                seq_index = processed_inputs.seq_index
-                seq_len = processed_inputs.seq_len
-
                 with get_train_context(
                     cp_size=self.cp_size,
                     cp_mesh=self.cp_mesh,
-                    cp_buffers=cp_buffers,
+                    cp_buffers=processed_inputs.cp_buffers,
                     sequence_dim=sequence_dim,
                     dtype=self.dtype,
                     autocast_enabled=self.autocast_enabled,
                 ):
-                    model_args = dict(
-                        input_ids=input_ids,
-                        attention_mask=attention_mask,
-                        position_ids=position_ids,
-                        use_cache=False,
+                    # Use forward_with_post_processing_fn for forward pass and post-processing
+                    rm_scores, _metrics, _ = forward_with_post_processing_fn(
+                        model=self.model,
+                        cfg=self.cfg,
+                        post_processing_fn=score_post_processor,
+                        processed_mb=processed_mb,
+                        is_reward_model=True,
+                        allow_flash_attn_args=False,
+                        sequence_dim=sequence_dim,
                     )
-                    outputs = self.model(**model_args)
-
-                    if not hasattr(outputs, "logits"):
-                        logits = self.model.lm_head(outputs.last_hidden_state)
-                    else:
-                        logits = outputs.logits
-                    # Apply temperature scaling
-                    logits = self._apply_temperature_scaling(logits)
-                if isinstance(logits, DTensor):
-                    logits = logits.to(torch.float32)
-                else:
-                    logits = outputs.logits.to(torch.float32)
-
-                rm_scores = to_local_if_dtensor(logits)
-                rm_scores = rm_scores.squeeze(-1)
 
                 # skip keeping the scores for the dummy batches
                 if batch_idx >= iterator_len:
@@ -990,13 +683,24 @@ def get_topk_logits(
             else self.cfg["logprob_batch_size"]
         )
 
-        sequence_dim = 1
-        seq_dim_size = data.get("input_ids").shape[sequence_dim]
+        # Validate sequence dimension
+        sequence_dim, seq_dim_size = check_sequence_dim(data)
 
         out_topk_vals = []
         out_topk_idx = []
         self.model.eval()
 
+        # Create top-k post-processor
+        topk_post_processor = TopkLogitsPostProcessor(
+            cfg=self.cfg,
+            device_mesh=self.device_mesh,
+            cp_mesh=self.cp_mesh,
+            tp_mesh=self.tp_mesh,
+            cp_size=self.cp_size,
+            k=k,
+            enable_seq_packing=self.enable_seq_packing,
+        )
+
         with torch.no_grad():
             data.to("cuda")
             # Get microbatch iterator based on batching strategy
@@ -1010,161 +714,27 @@ def get_topk_logits(
             )
 
             for batch_idx, processed_mb in enumerate(processed_iterator):
-                # Extract data dict and processed inputs
-                lp_batch = processed_mb.data_dict
                 processed_inputs = processed_mb.processed_inputs
-                input_lengths = lp_batch.get("input_lengths")
-
-                # Use original shapes from ProcessedMicrobatch (needed for unpacking later)
-                original_batch_size = processed_mb.original_batch_size
-                original_seq_len = processed_mb.original_seq_len
-
-                # Extract values from processed inputs
-                input_ids = processed_inputs.input_ids
-                attention_mask = processed_inputs.attention_mask
-                position_ids = processed_inputs.position_ids
-                flash_attn_kwargs = processed_inputs.flash_attn_kwargs
-                vlm_kwargs = processed_inputs.vlm_kwargs
-                cp_buffers = processed_inputs.cp_buffers
-                seq_index = processed_inputs.seq_index
-                seq_len = processed_inputs.seq_len
-                batch_size = input_ids.shape[0]
-
-                # Create all-ones attention mask for model input (required by DTensor)
-                with torch.autocast(device_type="cuda", dtype=self.dtype):
-                    attention_mask_input_all_ones = torch.ones(
-                        (batch_size, seq_len), dtype=torch.long, device=input_ids.device
-                    )
 
                 with get_train_context(
                     cp_size=self.cp_size,
                     cp_mesh=self.cp_mesh,
-                    cp_buffers=cp_buffers,
+                    cp_buffers=processed_inputs.cp_buffers,
                     sequence_dim=sequence_dim,
                     dtype=self.dtype,
                     autocast_enabled=self.autocast_enabled,
                 ):
-                    outputs = self.model(
-                        input_ids=input_ids,
-                        attention_mask=attention_mask_input_all_ones,
-                        position_ids=position_ids,
-                        use_cache=False,
-                        flash_attn_kwargs=flash_attn_kwargs,
-                    )
-
-                    if not hasattr(outputs, "logits"):
-                        logits = self.model.lm_head(outputs.last_hidden_state)
-                    else:
-                        logits = outputs.logits
-                    del outputs
-
-                    # Apply temperature scaling
-                    logits = self._apply_temperature_scaling(logits)
-
-                    if self.cp_size > 1:
-                        if isinstance(logits, DTensor):
-                            # Must be tp sharded
-                            assert (
-                                logits.device_mesh.ndim == 1
-                                and logits.device_mesh.mesh_dim_names[0] == "tp"
-                            ), "logits must be tp sharded"
-
-                            # CP is implicitly sharded on the seq dim, so we need to redistribute to the tp dim
-                            logits = DTensor.from_local(
-                                logits.to_local(),
-                                device_mesh=self.device_mesh[("cp", "tp")],
-                                placements=[Shard(sequence_dim), Shard(-1)],
-                            )
-                        else:
-                            logits = DTensor.from_local(
-                                logits,
-                                device_mesh=self.device_mesh[("cp", "tp")],
-                                placements=[Shard(sequence_dim), Shard(-1)],
-                            )
-
-                        # deal with TP first
-                        local_logits = logits.to_local()  # [B, S_cp, V_tp]
-
-                        tp_group = self.tp_mesh.get_group()
-                        tp_rank = torch.distributed.get_rank(tp_group)
-                        V_local = int(local_logits.shape[-1])
-                        vocab_start_index = tp_rank * V_local
-                        vocab_end_index = (tp_rank + 1) * V_local
-
-                        vals, idx = distributed_vocab_topk(
-                            local_logits,
-                            k=k,
-                            tp_group=tp_group,
-                            vocab_start_index=vocab_start_index,
-                            vocab_end_index=vocab_end_index,
-                        )
-                        # [B, S_cp, k]
-
-                        cp_group = self.cp_mesh.get_group()
-
-                        vals = allgather_cp_sharded_tensor(
-                            vals, cp_group, seq_dim=sequence_dim
-                        )
-                        idx = allgather_cp_sharded_tensor(
-                            idx, cp_group, seq_dim=sequence_dim
-                        )
-                        # [B, S, k]
-                    else:
-                        # Compute top-k over full sequence length (do not drop last position)
-                        if isinstance(logits, DTensor):
-                            local_logits = logits.to_local()  # [B, S, V_local]
-                            tp_group = self.tp_mesh.get_group()
-                            tp_rank = torch.distributed.get_rank(tp_group)
-                            V_local = int(local_logits.shape[-1])
-                            vocab_start_index = tp_rank * V_local
-                            vocab_end_index = (tp_rank + 1) * V_local
-
-                            vals, idx = distributed_vocab_topk(
-                                local_logits,
-                                k=k,
-                                tp_group=tp_group,
-                                vocab_start_index=vocab_start_index,
-                                vocab_end_index=vocab_end_index,
-                            )
-                        else:
-                            full_logits = logits.to(torch.float32)
-                            vals, idx = torch.topk(full_logits, k=k, dim=-1)
-
-                # Handle sequence packing unpacking
-                if self.enable_seq_packing:
-                    # Unpack top-k results from packed format back to original batch format
-                    # vals: [1, packed_seq_len, k] -> [original_batch_size, original_seq_len, k]
-                    # idx: [1, packed_seq_len, k] -> [original_batch_size, original_seq_len, k]
-
-                    # Create tensors to store unpacked results
-                    unpacked_vals = torch.zeros(
-                        (original_batch_size, original_seq_len, k),
-                        dtype=vals.dtype,
-                        device=vals.device,
-                    )
-                    unpacked_idx = torch.zeros(
-                        (original_batch_size, original_seq_len, k),
-                        dtype=idx.dtype,
-                        device=idx.device,
+                    # Use forward_with_post_processing_fn for forward pass and post-processing
+                    (vals, idx), _metrics, _ = forward_with_post_processing_fn(
+                        model=self.model,
+                        cfg=self.cfg,
+                        post_processing_fn=topk_post_processor,
+                        processed_mb=processed_mb,
+                        is_reward_model=False,
+                        allow_flash_attn_args=self.allow_flash_attn_args,
+                        sequence_dim=sequence_dim,
                     )
 
-                    # Get cumulative sequence lengths for unpacking
-                    cu_seqlens = flash_attn_kwargs.cu_seqlens_q
-
-                    for i in range(original_batch_size):
-                        start = cu_seqlens[i].item()
-                        end = cu_seqlens[i + 1].item()
-                        seq_len_actual = input_lengths[i].item()
-
-                        # Extract the corresponding portion from packed results
-                        # Note: vals and idx are [1, packed_seq_len, k] due to packing
-                        unpacked_vals[i, :seq_len_actual, :] = vals[0, start:end, :]
-                        unpacked_idx[i, :seq_len_actual, :] = idx[0, start:end, :]
-
-                    # Replace with unpacked results
-                    vals = unpacked_vals
-                    idx = unpacked_idx
-
                 # skip keeping the topk values for the dummy batches
                 if batch_idx >= iterator_len:
                     continue
diff --git a/tests/unit/models/automodel/test_automodel_train.py b/tests/unit/models/automodel/test_automodel_train.py
new file mode 100644
index 0000000000..1de03ef5bd
--- /dev/null
+++ b/tests/unit/models/automodel/test_automodel_train.py
@@ -0,0 +1,2075 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from torch.distributed.tensor import DTensor
+
+try:
+    import nemo_automodel  # noqa: F401
+except ImportError:
+    pytest.skip("nemo_automodel not available", allow_module_level=True)
+
+from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+from nemo_rl.models.automodel.data import (
+    ProcessedInputs,
+    ProcessedMicrobatch,
+    check_sequence_dim,
+    make_processed_microbatch_iterator,
+)
+from nemo_rl.models.automodel.train import (
+    LogprobsPostProcessor,
+    LossPostProcessor,
+    ScorePostProcessor,
+    TopkLogitsPostProcessor,
+    apply_temperature_scaling,
+    automodel_forward_backward,
+    extract_logits,
+    forward_with_post_processing_fn,
+    model_forward,
+)
+
+
+@pytest.fixture
+def mock_tokenizer():
+    tokenizer = MagicMock()
+    tokenizer.eos_token_id = 2
+    tokenizer.pad_token_id = 0
+    return tokenizer
+
+
+@pytest.fixture
+def mock_model():
+    model = MagicMock()
+    model.return_value = MagicMock(logits=torch.randn(4, 64, 32000))
+    return model
+
+
+@pytest.fixture
+def mock_loss_fn():
+    loss_fn = MagicMock()
+    loss_fn.return_value = (torch.tensor(0.5), {"loss": 0.5})
+    return loss_fn
+
+
+@pytest.fixture
+def mock_device_mesh():
+    mesh = MagicMock()
+    mesh.get_group.return_value = MagicMock()
+    mesh.__getitem__ = MagicMock(return_value=mesh)
+    return mesh
+
+
+@pytest.fixture
+def mock_cp_mesh():
+    mesh = MagicMock()
+    mesh.get_group.return_value = MagicMock()
+    return mesh
+
+
+@pytest.fixture
+def mock_tp_mesh():
+    mesh = MagicMock()
+    mesh.get_group.return_value = MagicMock()
+    return mesh
+
+
+@pytest.fixture
+def base_cfg():
+    return {
+        "dtensor_cfg": {"sequence_parallel": False},
+        "sequence_packing": {"train_mb_tokens": 256},
+        "generation": {"temperature": 1.0},
+    }
+
+
+@pytest.fixture
+def processed_inputs_no_flash():
+    return ProcessedInputs(
+        input_ids=torch.randint(0, 1000, (4, 64)),
+        seq_len=64,
+        attention_mask=torch.ones(4, 64, dtype=torch.bool),
+        position_ids=torch.arange(64).repeat(4, 1),
+        flash_attn_kwargs={},
+        vlm_kwargs={},
+        cp_buffers=[],
+        seq_index=None,
+    )
+
+
+@pytest.fixture
+def processed_inputs_with_flash():
+    @dataclass
+    class MockFlashAttnKwargs:
+        cu_seqlens_q: torch.Tensor
+
+    flash_kwargs = MockFlashAttnKwargs(cu_seqlens_q=torch.tensor([0, 32, 64, 96, 128]))
+    return ProcessedInputs(
+        input_ids=torch.randint(0, 1000, (1, 128)),
+        seq_len=128,
+        attention_mask=None,
+        position_ids=torch.arange(128).unsqueeze(0),
+        flash_attn_kwargs=flash_kwargs,
+        vlm_kwargs={},
+        cp_buffers=[],
+        seq_index=None,
+    )
+
+
+@pytest.fixture
+def processed_inputs_multimodal():
+    return ProcessedInputs(
+        input_ids=torch.randint(0, 1000, (2, 64)),
+        seq_len=64,
+        attention_mask=torch.ones(2, 64, dtype=torch.bool),
+        position_ids=None,
+        flash_attn_kwargs={},
+        vlm_kwargs={"pixel_values": torch.randn(2, 3, 224, 224)},
+        cp_buffers=[],
+        seq_index=None,
+    )
+
+
+# =====================
+# Test model_forward
+# =====================
+@pytest.mark.automodel
+class TestModelForward:
+    def test_basic_forward(self, mock_model, processed_inputs_no_flash):
+        result = model_forward(mock_model, processed_inputs_no_flash)
+
+        mock_model.assert_called_once()
+        call_kwargs = mock_model.call_args[1]
+        assert "input_ids" in call_kwargs
+        assert "attention_mask" in call_kwargs
+        assert "position_ids" in call_kwargs
+        assert call_kwargs["use_cache"] is False
+
+    def test_forward_with_flash_attention(
+        self, mock_model, processed_inputs_with_flash
+    ):
+        result = model_forward(mock_model, processed_inputs_with_flash)
+
+        mock_model.assert_called_once()
+        call_kwargs = mock_model.call_args[1]
+        assert "flash_attn_kwargs" in call_kwargs
+
+    def test_forward_with_multimodal(self, mock_model, processed_inputs_multimodal):
+        result = model_forward(mock_model, processed_inputs_multimodal)
+
+        mock_model.assert_called_once()
+        call_kwargs = mock_model.call_args[1]
+        assert "pixel_values" in call_kwargs
+        # Flash attention should be removed for multimodal
+        assert "flash_attn_kwargs" not in call_kwargs
+
+    def test_forward_reward_model_removes_flash_attn(
+        self, mock_model, processed_inputs_with_flash
+    ):
+        result = model_forward(
+            mock_model, processed_inputs_with_flash, is_reward_model=True
+        )
+
+        mock_model.assert_called_once()
+        call_kwargs = mock_model.call_args[1]
+        # Flash attention should be removed for reward models
+        assert "flash_attn_kwargs" not in call_kwargs
+
+    def test_forward_disallow_flash_attn_args(
+        self, mock_model, processed_inputs_with_flash
+    ):
+        result = model_forward(
+            mock_model, processed_inputs_with_flash, allow_flash_attn_args=False
+        )
+
+        mock_model.assert_called_once()
+        call_kwargs = mock_model.call_args[1]
+        assert "flash_attn_kwargs" not in call_kwargs
+
+
+# =====================
+# Test extract_logits
+# =====================
+@pytest.mark.automodel
+class TestExtractLogits:
+    def test_tensor_output(self, mock_model):
+        tensor_output = torch.randn(4, 64, 32000)
+        result = extract_logits(mock_model, tensor_output)
+        assert torch.equal(result, tensor_output)
+
+    def test_output_with_logits_attribute(self, mock_model):
+        mock_output = MagicMock()
+        mock_output.logits = torch.randn(4, 64, 32000)
+        result = extract_logits(mock_model, mock_output)
+        assert torch.equal(result, mock_output.logits)
+
+    def test_output_with_last_hidden_state(self, mock_model):
+        mock_output = MagicMock(spec=["last_hidden_state"])
+        mock_output.last_hidden_state = torch.randn(4, 64, 4096)
+        mock_model.lm_head = MagicMock(return_value=torch.randn(4, 64, 32000))
+
+        result = extract_logits(mock_model, mock_output)
+
+        mock_model.lm_head.assert_called_once_with(mock_output.last_hidden_state)
+
+
+# =====================
+# Test apply_temperature_scaling
+# =====================
+@pytest.mark.automodel
+class TestApplyTemperatureScaling:
+    def test_temperature_scaling_applied(self):
+        logits = torch.randn(4, 64, 32000)
+        original_logits = logits.clone()
+        cfg = {"generation": {"temperature": 2.0}}
+
+        result = apply_temperature_scaling(logits, cfg)
+
+        # Should be divided by temperature
+        expected = original_logits / 2.0
+        assert torch.allclose(result, expected)
+
+    def test_no_scaling_without_generation_config(self):
+        logits = torch.randn(4, 64, 32000)
+        original_logits = logits.clone()
+        cfg = {}
+
+        result = apply_temperature_scaling(logits, cfg)
+
+        assert torch.equal(result, original_logits)
+
+    def test_no_scaling_with_none_generation(self):
+        logits = torch.randn(4, 64, 32000)
+        original_logits = logits.clone()
+        cfg = {"generation": None}
+
+        result = apply_temperature_scaling(logits, cfg)
+
+        assert torch.equal(result, original_logits)
+
+
+# =====================
+# Test LossPostProcessor
+# =====================
+@pytest.mark.automodel
+class TestLossPostProcessor:
+    def test_basic_loss_computation(
+        self,
+        base_cfg,
+        mock_loss_fn,
+        mock_device_mesh,
+        mock_cp_mesh,
+        mock_tp_mesh,
+        processed_inputs_no_flash,
+    ):
+        processor = LossPostProcessor(
+            loss_fn=mock_loss_fn,
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+            dp_size=1,
+        )
+
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+
+        logits = torch.randn(batch_size, seq_len, vocab_size)
+        mb = BatchedDataDict(
+            {
+                "input_ids": torch.randint(0, vocab_size, (batch_size, seq_len)),
+                "sample_mask": torch.ones(batch_size, dtype=torch.bool),
+            }
+        )
+        global_valid_seqs = torch.tensor(8)
+        global_valid_toks = torch.tensor(512)
+
+        loss, metrics = processor(
+            logits=logits,
+            mb=mb,
+            processed_inputs=processed_inputs_no_flash,
+            global_valid_seqs=global_valid_seqs,
+            global_valid_toks=global_valid_toks,
+        )
+
+        # Verify loss function was called
+        mock_loss_fn.assert_called_once()
+        call_args = mock_loss_fn.call_args[0]
+        assert torch.is_tensor(call_args[0])  # logits
+        assert call_args[2] == global_valid_seqs  # global_valid_seqs
+        assert call_args[3] == global_valid_toks  # global_valid_toks
+
+    @patch("nemo_rl.models.automodel.train.SequencePackingLossWrapper")
+    def test_loss_with_sequence_packing(
+        self,
+        mock_wrapper_class,
+        base_cfg,
+        mock_loss_fn,
+        mock_device_mesh,
+        mock_cp_mesh,
+        mock_tp_mesh,
+        processed_inputs_with_flash,
+    ):
+        # Setup mock wrapper
+        mock_wrapper_instance = MagicMock()
+        mock_wrapper_instance.return_value = (torch.tensor(0.5), {"loss": 0.5})
+        mock_wrapper_class.return_value = mock_wrapper_instance
+
+        processor = LossPostProcessor(
+            loss_fn=mock_loss_fn,
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+            dp_size=1,
+        )
+
+        batch_size = 1
+        seq_len = 128
+        vocab_size = 32000
+
+        logits = torch.randn(batch_size, seq_len, vocab_size)
+        mb = BatchedDataDict(
+            {
+                "input_ids": torch.randint(0, vocab_size, (batch_size, seq_len)),
+                "sample_mask": torch.ones(batch_size, dtype=torch.bool),
+            }
+        )
+        global_valid_seqs = torch.tensor(4)
+        global_valid_toks = torch.tensor(128)
+
+        loss, metrics = processor(
+            logits=logits,
+            mb=mb,
+            processed_inputs=processed_inputs_with_flash,
+            global_valid_seqs=global_valid_seqs,
+            global_valid_toks=global_valid_toks,
+        )
+
+        # Verify SequencePackingLossWrapper was created
+        mock_wrapper_class.assert_called_once()
+        # Verify the wrapper was called instead of raw loss_fn
+        mock_wrapper_instance.assert_called_once()
+
+    def test_loss_processor_initialization(
+        self,
+        base_cfg,
+        mock_loss_fn,
+        mock_device_mesh,
+        mock_cp_mesh,
+        mock_tp_mesh,
+    ):
+        processor = LossPostProcessor(
+            loss_fn=mock_loss_fn,
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=2,
+            dp_size=4,
+        )
+
+        assert processor.loss_fn is mock_loss_fn
+        assert processor.cfg is base_cfg
+        assert processor.cp_size == 2
+        assert processor.dp_size == 4
+
+
+# =====================
+# Test ScorePostProcessor
+# =====================
+@pytest.mark.automodel
+class TestScorePostProcessor:
+    def test_basic_scoring(self, base_cfg):
+        processor = ScorePostProcessor(cfg=base_cfg)
+
+        # Create mock logits with shape [batch_size, 1]
+        logits = torch.randn(4, 1)
+
+        result = processor(logits)
+
+        assert result.shape == (4,)
+        assert result.dtype == torch.float32
+
+    def test_scoring_with_sequence_logits(self, base_cfg):
+        processor = ScorePostProcessor(cfg=base_cfg)
+
+        # Create mock logits with shape [batch_size, seq_len, 1]
+        logits = torch.randn(4, 64, 1)
+
+        result = processor(logits)
+
+        assert result.shape == (4, 64)
+        assert result.dtype == torch.float32
+
+
+# =====================
+# Test LogprobsPostProcessor
+# =====================
+@pytest.mark.automodel
+class TestLogprobsPostProcessor:
+    def test_basic_logprobs_computation(
+        self, base_cfg, mock_device_mesh, mock_cp_mesh, mock_tp_mesh
+    ):
+        processor = LogprobsPostProcessor(
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+        )
+
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+
+        logits = torch.randn(batch_size, seq_len, vocab_size)
+        input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+        input_lengths = torch.full((batch_size,), seq_len)
+
+        processed_inputs = ProcessedInputs(
+            input_ids=input_ids,
+            seq_len=seq_len,
+            attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+            position_ids=torch.arange(seq_len).repeat(batch_size, 1),
+            flash_attn_kwargs={},
+            vlm_kwargs={},
+            cp_buffers=[],
+            seq_index=None,
+        )
+
+        result = processor(
+            logits=logits,
+            processed_inputs=processed_inputs,
+            input_lengths=input_lengths,
+            original_batch_size=batch_size,
+            original_seq_len=seq_len,
+            enable_seq_packing=False,
+        )
+
+        assert result.shape == (batch_size, seq_len)
+
+    def test_logprobs_with_chunking(
+        self, base_cfg, mock_device_mesh, mock_cp_mesh, mock_tp_mesh
+    ):
+        cfg_with_chunk = {**base_cfg, "logprob_chunk_size": 16}
+        processor = LogprobsPostProcessor(
+            cfg=cfg_with_chunk,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+        )
+
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+
+        logits = torch.randn(batch_size, seq_len, vocab_size)
+        input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+        input_lengths = torch.full((batch_size,), seq_len)
+
+        processed_inputs = ProcessedInputs(
+            input_ids=input_ids,
+            seq_len=seq_len,
+            attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+            position_ids=torch.arange(seq_len).repeat(batch_size, 1),
+            flash_attn_kwargs={},
+            vlm_kwargs={},
+            cp_buffers=[],
+            seq_index=None,
+        )
+
+        result = processor(
+            logits=logits,
+            processed_inputs=processed_inputs,
+            input_lengths=input_lengths,
+            original_batch_size=batch_size,
+            original_seq_len=seq_len,
+            enable_seq_packing=False,
+        )
+
+        assert result.shape == (batch_size, seq_len)
+
+
+# =====================
+# Test TopkLogitsPostProcessor
+# =====================
+@pytest.mark.automodel
+class TestTopkLogitsPostProcessor:
+    def test_basic_topk(self, base_cfg, mock_device_mesh, mock_cp_mesh, mock_tp_mesh):
+        k = 10
+        processor = TopkLogitsPostProcessor(
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+            k=k,
+        )
+
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+
+        logits = torch.randn(batch_size, seq_len, vocab_size)
+        input_lengths = torch.full((batch_size,), seq_len)
+
+        processed_inputs = ProcessedInputs(
+            input_ids=torch.randint(0, vocab_size, (batch_size, seq_len)),
+            seq_len=seq_len,
+            attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+            position_ids=torch.arange(seq_len).repeat(batch_size, 1),
+            flash_attn_kwargs={},
+            vlm_kwargs={},
+            cp_buffers=[],
+            seq_index=None,
+        )
+
+        vals, idx = processor(
+            logits=logits,
+            processed_inputs=processed_inputs,
+            input_lengths=input_lengths,
+            original_batch_size=batch_size,
+            original_seq_len=seq_len,
+            enable_seq_packing=False,
+        )
+
+        assert vals.shape == (batch_size, seq_len, k)
+        assert idx.shape == (batch_size, seq_len, k)
+
+
+# =====================
+# Test ProcessedMicrobatch
+# =====================
+@pytest.mark.automodel
+class TestProcessedMicrobatch:
+    def test_processed_microbatch_creation(self, processed_inputs_no_flash):
+        data_dict = BatchedDataDict(
+            {
+                "input_ids": torch.randint(0, 1000, (4, 64)),
+                "sample_mask": torch.ones(4, dtype=torch.bool),
+            }
+        )
+
+        pm = ProcessedMicrobatch(
+            data_dict=data_dict,
+            processed_inputs=processed_inputs_no_flash,
+            original_batch_size=4,
+            original_seq_len=64,
+        )
+
+        assert pm.original_batch_size == 4
+        assert pm.original_seq_len == 64
+        assert pm.data_dict is data_dict
+        assert pm.processed_inputs is processed_inputs_no_flash
+
+
+# =====================
+# Test make_processed_microbatch_iterator
+# =====================
+@pytest.mark.automodel
+class TestMakeProcessedMicrobatchIterator:
+    def test_basic_iteration(self, mock_tokenizer):
+        # Create test data
+        data_dict1 = BatchedDataDict(
+            {
+                "input_ids": torch.randint(0, 1000, (4, 64)),
+                "sample_mask": torch.ones(4, dtype=torch.bool),
+            }
+        )
+        data_dict2 = BatchedDataDict(
+            {
+                "input_ids": torch.randint(0, 1000, (4, 64)),
+                "sample_mask": torch.ones(4, dtype=torch.bool),
+            }
+        )
+
+        # Mock get_multimodal_dict to return empty dict
+        data_dict1.get_multimodal_dict = MagicMock(return_value={})
+        data_dict2.get_multimodal_dict = MagicMock(return_value={})
+
+        raw_iterator = iter([data_dict1, data_dict2])
+
+        cfg = {
+            "dtensor_cfg": {"sequence_parallel": False},
+            "sequence_packing": {"enabled": False},
+        }
+
+        processed_iterator = make_processed_microbatch_iterator(
+            raw_iterator=raw_iterator,
+            tokenizer=mock_tokenizer,
+            cfg=cfg,
+            cp_size=1,
+        )
+
+        results = list(processed_iterator)
+
+        assert len(results) == 2
+        assert all(isinstance(pm, ProcessedMicrobatch) for pm in results)
+        assert all(pm.original_batch_size == 4 for pm in results)
+        assert all(pm.original_seq_len == 64 for pm in results)
+
+
+# =====================
+# Test check_sequence_dim
+# =====================
+@pytest.mark.automodel
+class TestCheckSequenceDim:
+    def test_consistent_sequence_dim(self):
+        data = BatchedDataDict(
+            {
+                "input_ids": torch.randint(0, 1000, (4, 64)),
+                "attention_mask": torch.ones(4, 64, dtype=torch.bool),
+                "token_mask": torch.ones(4, 64, dtype=torch.bool),
+                "sample_mask": torch.ones(4, dtype=torch.bool),  # 1D tensor
+            }
+        )
+
+        seq_dim, seq_dim_size = check_sequence_dim(data)
+
+        assert seq_dim == 1
+        assert seq_dim_size == 64
+
+    def test_inconsistent_sequence_dim_raises_error(self):
+        data = BatchedDataDict(
+            {
+                "input_ids": torch.randint(0, 1000, (4, 64)),
+                "attention_mask": torch.ones(
+                    4, 128, dtype=torch.bool
+                ),  # Different seq len
+            }
+        )
+
+        with pytest.raises(AssertionError, match="Dim 1 must be the sequence dim"):
+            check_sequence_dim(data)
+
+    def test_ignores_1d_tensors(self):
+        data = BatchedDataDict(
+            {
+                "input_ids": torch.randint(0, 1000, (4, 64)),
+                "sample_mask": torch.ones(4, dtype=torch.bool),  # 1D tensor
+                "labels": torch.randint(0, 2, (128,)),  # Different 1D tensor size
+            }
+        )
+
+        seq_dim, seq_dim_size = check_sequence_dim(data)
+
+        assert seq_dim == 1
+        assert seq_dim_size == 64
+
+
+# =====================
+# Test ProcessedInputs properties
+# =====================
+@pytest.mark.automodel
+class TestProcessedInputsProperties:
+    def test_has_context_parallel_false(self, processed_inputs_no_flash):
+        assert processed_inputs_no_flash.has_context_parallel is False
+
+    def test_has_context_parallel_true(self):
+        input_ids = torch.randint(0, 1000, (2, 128))
+        position_ids = torch.arange(128).repeat(2, 1)
+        seq_index = torch.arange(128).repeat(1, 1)
+
+        processed = ProcessedInputs(
+            input_ids=input_ids,
+            seq_len=128,
+            attention_mask=torch.ones(2, 128, dtype=torch.bool),
+            position_ids=position_ids,
+            flash_attn_kwargs={},
+            vlm_kwargs={},
+            cp_buffers=[input_ids, position_ids, seq_index],
+            seq_index=seq_index,
+        )
+
+        assert processed.has_context_parallel is True
+
+    def test_has_flash_attention_false(self, processed_inputs_no_flash):
+        assert processed_inputs_no_flash.has_flash_attention is False
+
+    def test_has_flash_attention_true(self, processed_inputs_with_flash):
+        assert processed_inputs_with_flash.has_flash_attention is True
+
+    def test_is_multimodal_false(self, processed_inputs_no_flash):
+        assert processed_inputs_no_flash.is_multimodal is False
+
+    def test_is_multimodal_true(self, processed_inputs_multimodal):
+        assert processed_inputs_multimodal.is_multimodal is True
+
+
+# =====================
+# Test forward_with_post_processing_fn
+# =====================
+@pytest.mark.automodel
+class TestForwardWithPostProcessingFn:
+    def test_forward_with_loss_post_processor(
+        self,
+        mock_model,
+        mock_loss_fn,
+        base_cfg,
+        mock_device_mesh,
+        mock_cp_mesh,
+        mock_tp_mesh,
+    ):
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+
+        # Create processed inputs
+        input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+        processed_inputs = ProcessedInputs(
+            input_ids=input_ids,
+            seq_len=seq_len,
+            attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+            position_ids=torch.arange(seq_len).repeat(batch_size, 1),
+            flash_attn_kwargs={},
+            vlm_kwargs={},
+            cp_buffers=[],
+            seq_index=None,
+        )
+
+        # Create data dict
+        data_dict = BatchedDataDict(
+            {
+                "input_ids": input_ids,
+                "input_lengths": torch.full((batch_size,), seq_len),
+                "sample_mask": torch.ones(batch_size, dtype=torch.bool),
+            }
+        )
+
+        # Create processed microbatch
+        processed_mb = ProcessedMicrobatch(
+            data_dict=data_dict,
+            processed_inputs=processed_inputs,
+            original_batch_size=batch_size,
+            original_seq_len=seq_len,
+        )
+
+        # Create loss post-processor
+        loss_post_processor = LossPostProcessor(
+            loss_fn=mock_loss_fn,
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+            dp_size=1,
+        )
+
+        # Call forward_with_post_processing_fn
+        result, metrics, returned_mb = forward_with_post_processing_fn(
+            model=mock_model,
+            cfg=base_cfg,
+            post_processing_fn=loss_post_processor,
+            processed_mb=processed_mb,
+            global_valid_seqs=torch.tensor(batch_size),
+            global_valid_toks=torch.tensor(batch_size * seq_len),
+        )
+
+        # Verify model was called
+        mock_model.assert_called_once()
+
+        # Verify loss function was called
+        mock_loss_fn.assert_called_once()
+
+        # Verify returned microbatch is correct
+        assert returned_mb is processed_mb
+
+    def test_forward_with_score_post_processor(
+        self,
+        mock_model,
+        base_cfg,
+    ):
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+
+        # Setup mock model to return reward-like logits
+        mock_model.return_value = MagicMock(logits=torch.randn(batch_size, 1))
+
+        # Create processed inputs
+        input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+        processed_inputs = ProcessedInputs(
+            input_ids=input_ids,
+            seq_len=seq_len,
+            attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+            position_ids=torch.arange(seq_len).repeat(batch_size, 1),
+            flash_attn_kwargs={},
+            vlm_kwargs={},
+            cp_buffers=[],
+            seq_index=None,
+        )
+
+        # Create data dict
+        data_dict = BatchedDataDict(
+            {
+                "input_ids": input_ids,
+                "input_lengths": torch.full((batch_size,), seq_len),
+                "sample_mask": torch.ones(batch_size, dtype=torch.bool),
+            }
+        )
+
+        # Create processed microbatch
+        processed_mb = ProcessedMicrobatch(
+            data_dict=data_dict,
+            processed_inputs=processed_inputs,
+            original_batch_size=batch_size,
+            original_seq_len=seq_len,
+        )
+
+        # Create score post-processor
+        score_post_processor = ScorePostProcessor(cfg=base_cfg)
+
+        # Call forward_with_post_processing_fn
+        result, metrics, returned_mb = forward_with_post_processing_fn(
+            model=mock_model,
+            cfg=base_cfg,
+            post_processing_fn=score_post_processor,
+            processed_mb=processed_mb,
+            is_reward_model=True,
+        )
+
+        # Verify model was called
+        mock_model.assert_called_once()
+
+        # Verify scores are in metrics
+        assert "scores" in metrics
+
+        # Verify result shape
+        assert result.shape == (batch_size,)
+
+
+# =====================
+# Test automodel_forward_backward
+# =====================
+@pytest.mark.automodel
+class TestAutomodelForwardBackward:
+    def test_forward_backward_single_microbatch(
+        self,
+        mock_model,
+        mock_loss_fn,
+        base_cfg,
+        mock_device_mesh,
+        mock_cp_mesh,
+        mock_tp_mesh,
+    ):
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+
+        # Create processed inputs
+        input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+        processed_inputs = ProcessedInputs(
+            input_ids=input_ids,
+            seq_len=seq_len,
+            attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+            position_ids=torch.arange(seq_len).repeat(batch_size, 1),
+            flash_attn_kwargs={},
+            vlm_kwargs={},
+            cp_buffers=[],
+            seq_index=None,
+        )
+
+        # Create data dict
+        data_dict = BatchedDataDict(
+            {
+                "input_ids": input_ids,
+                "input_lengths": torch.full((batch_size,), seq_len),
+                "sample_mask": torch.ones(batch_size, dtype=torch.bool),
+            }
+        )
+
+        # Create processed microbatch
+        processed_mb = ProcessedMicrobatch(
+            data_dict=data_dict,
+            processed_inputs=processed_inputs,
+            original_batch_size=batch_size,
+            original_seq_len=seq_len,
+        )
+
+        # Create loss post-processor
+        loss_post_processor = LossPostProcessor(
+            loss_fn=mock_loss_fn,
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+            dp_size=1,
+        )
+
+        # Call automodel_forward_backward in forward_only mode
+        results = automodel_forward_backward(
+            model=mock_model,
+            cfg=base_cfg,
+            data_iterator=iter([processed_mb]),
+            post_processing_fn=loss_post_processor,
+            forward_only=True,
+            global_valid_seqs=torch.tensor(batch_size),
+            global_valid_toks=torch.tensor(batch_size * seq_len),
+        )
+
+        # Verify results
+        assert len(results) == 1
+        _result, _metrics = results[0]
+
+        # Verify loss function was called
+        mock_loss_fn.assert_called_once()
+
+    def test_forward_backward_multiple_microbatches(
+        self,
+        mock_model,
+        mock_loss_fn,
+        base_cfg,
+        mock_device_mesh,
+        mock_cp_mesh,
+        mock_tp_mesh,
+    ):
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+        num_microbatches = 3
+
+        # Create multiple processed microbatches
+        processed_mbs = []
+        for _ in range(num_microbatches):
+            input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+            processed_inputs = ProcessedInputs(
+                input_ids=input_ids,
+                seq_len=seq_len,
+                attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+                position_ids=torch.arange(seq_len).repeat(batch_size, 1),
+                flash_attn_kwargs={},
+                vlm_kwargs={},
+                cp_buffers=[],
+                seq_index=None,
+            )
+
+            data_dict = BatchedDataDict(
+                {
+                    "input_ids": input_ids,
+                    "input_lengths": torch.full((batch_size,), seq_len),
+                    "sample_mask": torch.ones(batch_size, dtype=torch.bool),
+                }
+            )
+
+            processed_mbs.append(
+                ProcessedMicrobatch(
+                    data_dict=data_dict,
+                    processed_inputs=processed_inputs,
+                    original_batch_size=batch_size,
+                    original_seq_len=seq_len,
+                )
+            )
+
+        # Create loss post-processor
+        loss_post_processor = LossPostProcessor(
+            loss_fn=mock_loss_fn,
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+            dp_size=1,
+        )
+
+        # Call automodel_forward_backward in forward_only mode
+        results = automodel_forward_backward(
+            model=mock_model,
+            cfg=base_cfg,
+            data_iterator=iter(processed_mbs),
+            post_processing_fn=loss_post_processor,
+            forward_only=True,
+            global_valid_seqs=torch.tensor(batch_size * num_microbatches),
+            global_valid_toks=torch.tensor(batch_size * seq_len * num_microbatches),
+        )
+
+        # Verify results
+        assert len(results) == num_microbatches
+
+        # Verify model was called num_microbatches times
+        assert mock_model.call_count == num_microbatches
+
+        # Verify loss function was called num_microbatches times
+        assert mock_loss_fn.call_count == num_microbatches
+
+    def test_forward_backward_with_train_context_fn(
+        self,
+        mock_model,
+        mock_loss_fn,
+        base_cfg,
+        mock_device_mesh,
+        mock_cp_mesh,
+        mock_tp_mesh,
+    ):
+        """Test automodel_forward_backward with train_context_fn callback."""
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+
+        # Create processed inputs
+        input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+        processed_inputs = ProcessedInputs(
+            input_ids=input_ids,
+            seq_len=seq_len,
+            attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+            position_ids=torch.arange(seq_len).repeat(batch_size, 1),
+            flash_attn_kwargs={},
+            vlm_kwargs={},
+            cp_buffers=[],
+            seq_index=None,
+        )
+
+        # Create data dict
+        data_dict = BatchedDataDict(
+            {
+                "input_ids": input_ids,
+                "input_lengths": torch.full((batch_size,), seq_len),
+                "sample_mask": torch.ones(batch_size, dtype=torch.bool),
+            }
+        )
+
+        # Create processed microbatch
+        processed_mb = ProcessedMicrobatch(
+            data_dict=data_dict,
+            processed_inputs=processed_inputs,
+            original_batch_size=batch_size,
+            original_seq_len=seq_len,
+        )
+
+        # Create loss post-processor
+        loss_post_processor = LossPostProcessor(
+            loss_fn=mock_loss_fn,
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+            dp_size=1,
+        )
+
+        # Track context manager calls
+        context_calls = []
+
+        class MockContext:
+            def __enter__(self):
+                context_calls.append("enter")
+                return self
+
+            def __exit__(self, *args):
+                context_calls.append("exit")
+                return False
+
+        def mock_train_context_fn(processed_inputs):
+            return MockContext()
+
+        # Call automodel_forward_backward with train_context_fn
+        results = automodel_forward_backward(
+            model=mock_model,
+            cfg=base_cfg,
+            data_iterator=iter([processed_mb]),
+            post_processing_fn=loss_post_processor,
+            forward_only=True,
+            global_valid_seqs=torch.tensor(batch_size),
+            global_valid_toks=torch.tensor(batch_size * seq_len),
+            train_context_fn=mock_train_context_fn,
+        )
+
+        # Verify context manager was called
+        assert context_calls == ["enter", "exit"]
+        assert len(results) == 1
+
+    def test_forward_backward_with_on_microbatch_start(
+        self,
+        mock_model,
+        mock_loss_fn,
+        base_cfg,
+        mock_device_mesh,
+        mock_cp_mesh,
+        mock_tp_mesh,
+    ):
+        """Test automodel_forward_backward with on_microbatch_start callback."""
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+        num_microbatches = 3
+
+        # Create multiple processed microbatches
+        processed_mbs = []
+        for _ in range(num_microbatches):
+            input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+            processed_inputs = ProcessedInputs(
+                input_ids=input_ids,
+                seq_len=seq_len,
+                attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+                position_ids=torch.arange(seq_len).repeat(batch_size, 1),
+                flash_attn_kwargs={},
+                vlm_kwargs={},
+                cp_buffers=[],
+                seq_index=None,
+            )
+
+            data_dict = BatchedDataDict(
+                {
+                    "input_ids": input_ids,
+                    "input_lengths": torch.full((batch_size,), seq_len),
+                    "sample_mask": torch.ones(batch_size, dtype=torch.bool),
+                }
+            )
+
+            processed_mbs.append(
+                ProcessedMicrobatch(
+                    data_dict=data_dict,
+                    processed_inputs=processed_inputs,
+                    original_batch_size=batch_size,
+                    original_seq_len=seq_len,
+                )
+            )
+
+        # Create loss post-processor
+        loss_post_processor = LossPostProcessor(
+            loss_fn=mock_loss_fn,
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+            dp_size=1,
+        )
+
+        # Track callback calls
+        callback_indices = []
+
+        def on_microbatch_start(mb_idx):
+            callback_indices.append(mb_idx)
+
+        # Call automodel_forward_backward with on_microbatch_start
+        results = automodel_forward_backward(
+            model=mock_model,
+            cfg=base_cfg,
+            data_iterator=iter(processed_mbs),
+            post_processing_fn=loss_post_processor,
+            forward_only=True,
+            global_valid_seqs=torch.tensor(batch_size * num_microbatches),
+            global_valid_toks=torch.tensor(batch_size * seq_len * num_microbatches),
+            on_microbatch_start=on_microbatch_start,
+        )
+
+        # Verify callback was called for each microbatch
+        assert callback_indices == [0, 1, 2]
+        assert len(results) == num_microbatches
+
+    def test_forward_backward_with_dummy_batches(
+        self,
+        mock_model,
+        mock_loss_fn,
+        base_cfg,
+        mock_device_mesh,
+        mock_cp_mesh,
+        mock_tp_mesh,
+    ):
+        """Test automodel_forward_backward with dummy batches (num_valid_microbatches)."""
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+        total_microbatches = 3
+        num_valid_microbatches = 2  # Only first 2 are valid
+
+        # Create processed microbatches
+        processed_mbs = []
+        for _ in range(total_microbatches):
+            input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+            processed_inputs = ProcessedInputs(
+                input_ids=input_ids,
+                seq_len=seq_len,
+                attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+                position_ids=torch.arange(seq_len).repeat(batch_size, 1),
+                flash_attn_kwargs={},
+                vlm_kwargs={},
+                cp_buffers=[],
+                seq_index=None,
+            )
+
+            data_dict = BatchedDataDict(
+                {
+                    "input_ids": input_ids,
+                    "input_lengths": torch.full((batch_size,), seq_len),
+                    "sample_mask": torch.ones(batch_size, dtype=torch.bool),
+                }
+            )
+
+            processed_mbs.append(
+                ProcessedMicrobatch(
+                    data_dict=data_dict,
+                    processed_inputs=processed_inputs,
+                    original_batch_size=batch_size,
+                    original_seq_len=seq_len,
+                )
+            )
+
+        # Create loss post-processor
+        loss_post_processor = LossPostProcessor(
+            loss_fn=mock_loss_fn,
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+            dp_size=1,
+        )
+
+        # Call automodel_forward_backward with num_valid_microbatches
+        results = automodel_forward_backward(
+            model=mock_model,
+            cfg=base_cfg,
+            data_iterator=iter(processed_mbs),
+            post_processing_fn=loss_post_processor,
+            forward_only=True,
+            global_valid_seqs=torch.tensor(batch_size * num_valid_microbatches),
+            global_valid_toks=torch.tensor(
+                batch_size * seq_len * num_valid_microbatches
+            ),
+            num_valid_microbatches=num_valid_microbatches,
+        )
+
+        # Verify all microbatches processed
+        assert len(results) == total_microbatches
+
+        # Third batch (index 2) is dummy - result should be zeroed
+        dummy_result, dummy_metrics = results[2]
+        assert dummy_result.item() == 0.0  # Dummy batch loss is zeroed
+
+
+# =====================
+# Test forward_with_post_processing_fn (additional coverage)
+# =====================
+@pytest.mark.automodel
+class TestForwardWithPostProcessingFnAdditional:
+    def test_forward_with_logprobs_post_processor(
+        self,
+        mock_model,
+        base_cfg,
+        mock_device_mesh,
+        mock_cp_mesh,
+        mock_tp_mesh,
+    ):
+        """Test forward_with_post_processing_fn with LogprobsPostProcessor."""
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+
+        # Create processed inputs
+        input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+        processed_inputs = ProcessedInputs(
+            input_ids=input_ids,
+            seq_len=seq_len,
+            attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+            position_ids=torch.arange(seq_len).repeat(batch_size, 1),
+            flash_attn_kwargs={},
+            vlm_kwargs={},
+            cp_buffers=[],
+            seq_index=None,
+        )
+
+        # Create data dict
+        data_dict = BatchedDataDict(
+            {
+                "input_ids": input_ids,
+                "input_lengths": torch.full((batch_size,), seq_len),
+                "sample_mask": torch.ones(batch_size, dtype=torch.bool),
+            }
+        )
+
+        # Create processed microbatch
+        processed_mb = ProcessedMicrobatch(
+            data_dict=data_dict,
+            processed_inputs=processed_inputs,
+            original_batch_size=batch_size,
+            original_seq_len=seq_len,
+        )
+
+        # Create logprobs post-processor
+        logprobs_post_processor = LogprobsPostProcessor(
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+        )
+
+        # Call forward_with_post_processing_fn
+        result, metrics, returned_mb = forward_with_post_processing_fn(
+            model=mock_model,
+            cfg=base_cfg,
+            post_processing_fn=logprobs_post_processor,
+            processed_mb=processed_mb,
+        )
+
+        # Verify model was called
+        mock_model.assert_called_once()
+
+        # Verify logprobs are in metrics
+        assert "logprobs" in metrics
+
+        # Verify result shape
+        assert result.shape == (batch_size, seq_len)
+
+        # Verify returned microbatch is correct
+        assert returned_mb is processed_mb
+
+    def test_forward_with_topk_post_processor(
+        self,
+        mock_model,
+        base_cfg,
+        mock_device_mesh,
+        mock_cp_mesh,
+        mock_tp_mesh,
+    ):
+        """Test forward_with_post_processing_fn with TopkLogitsPostProcessor."""
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+        k = 10
+
+        # Create processed inputs
+        input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+        processed_inputs = ProcessedInputs(
+            input_ids=input_ids,
+            seq_len=seq_len,
+            attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+            position_ids=torch.arange(seq_len).repeat(batch_size, 1),
+            flash_attn_kwargs={},
+            vlm_kwargs={},
+            cp_buffers=[],
+            seq_index=None,
+        )
+
+        # Create data dict
+        data_dict = BatchedDataDict(
+            {
+                "input_ids": input_ids,
+                "input_lengths": torch.full((batch_size,), seq_len),
+                "sample_mask": torch.ones(batch_size, dtype=torch.bool),
+            }
+        )
+
+        # Create processed microbatch
+        processed_mb = ProcessedMicrobatch(
+            data_dict=data_dict,
+            processed_inputs=processed_inputs,
+            original_batch_size=batch_size,
+            original_seq_len=seq_len,
+        )
+
+        # Create topk post-processor
+        topk_post_processor = TopkLogitsPostProcessor(
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+            k=k,
+        )
+
+        # Call forward_with_post_processing_fn
+        result, metrics, returned_mb = forward_with_post_processing_fn(
+            model=mock_model,
+            cfg=base_cfg,
+            post_processing_fn=topk_post_processor,
+            processed_mb=processed_mb,
+        )
+
+        # Verify model was called
+        mock_model.assert_called_once()
+
+        # Verify topk values and indices are in metrics
+        assert "topk_logits" in metrics
+        assert "topk_indices" in metrics
+
+        # Verify result is tuple of (vals, idx)
+        vals, idx = result
+        assert vals.shape == (batch_size, seq_len, k)
+        assert idx.shape == (batch_size, seq_len, k)
+
+    def test_forward_with_unknown_post_processor_raises_error(
+        self,
+        mock_model,
+        base_cfg,
+    ):
+        """Test forward_with_post_processing_fn raises TypeError for unknown post-processor."""
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+
+        # Create processed inputs
+        input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+        processed_inputs = ProcessedInputs(
+            input_ids=input_ids,
+            seq_len=seq_len,
+            attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+            position_ids=torch.arange(seq_len).repeat(batch_size, 1),
+            flash_attn_kwargs={},
+            vlm_kwargs={},
+            cp_buffers=[],
+            seq_index=None,
+        )
+
+        # Create data dict
+        data_dict = BatchedDataDict(
+            {
+                "input_ids": input_ids,
+                "input_lengths": torch.full((batch_size,), seq_len),
+                "sample_mask": torch.ones(batch_size, dtype=torch.bool),
+            }
+        )
+
+        # Create processed microbatch
+        processed_mb = ProcessedMicrobatch(
+            data_dict=data_dict,
+            processed_inputs=processed_inputs,
+            original_batch_size=batch_size,
+            original_seq_len=seq_len,
+        )
+
+        # Create unknown post-processor (not a valid type)
+        class UnknownPostProcessor:
+            pass
+
+        unknown_post_processor = UnknownPostProcessor()
+
+        # Call forward_with_post_processing_fn and expect TypeError
+        with pytest.raises(TypeError, match="Unknown post-processing function type"):
+            forward_with_post_processing_fn(
+                model=mock_model,
+                cfg=base_cfg,
+                post_processing_fn=unknown_post_processor,
+                processed_mb=processed_mb,
+            )
+
+    def test_forward_with_processed_mb_directly(
+        self,
+        mock_model,
+        mock_loss_fn,
+        base_cfg,
+        mock_device_mesh,
+        mock_cp_mesh,
+        mock_tp_mesh,
+    ):
+        """Test forward_with_post_processing_fn with processed_mb provided directly."""
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+
+        # Create processed inputs
+        input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+        processed_inputs = ProcessedInputs(
+            input_ids=input_ids,
+            seq_len=seq_len,
+            attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+            position_ids=torch.arange(seq_len).repeat(batch_size, 1),
+            flash_attn_kwargs={},
+            vlm_kwargs={},
+            cp_buffers=[],
+            seq_index=None,
+        )
+
+        # Create data dict
+        data_dict = BatchedDataDict(
+            {
+                "input_ids": input_ids,
+                "input_lengths": torch.full((batch_size,), seq_len),
+                "sample_mask": torch.ones(batch_size, dtype=torch.bool),
+            }
+        )
+
+        # Create processed microbatch
+        processed_mb = ProcessedMicrobatch(
+            data_dict=data_dict,
+            processed_inputs=processed_inputs,
+            original_batch_size=batch_size,
+            original_seq_len=seq_len,
+        )
+
+        # Create loss post-processor
+        loss_post_processor = LossPostProcessor(
+            loss_fn=mock_loss_fn,
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+            dp_size=1,
+        )
+
+        # Call forward_with_post_processing_fn with processed_mb directly (no iterator)
+        result, metrics, returned_mb = forward_with_post_processing_fn(
+            model=mock_model,
+            cfg=base_cfg,
+            post_processing_fn=loss_post_processor,
+            processed_mb=processed_mb,  # Directly provided
+            global_valid_seqs=torch.tensor(batch_size),
+            global_valid_toks=torch.tensor(batch_size * seq_len),
+        )
+
+        # Verify model was called
+        mock_model.assert_called_once()
+
+        # Verify returned microbatch is correct
+        assert returned_mb is processed_mb
+
+
+# =====================
+# Test redistribute_logits_for_cp
+# =====================
+@pytest.mark.automodel
+class TestRedistributeLogitsForCP:
+    def test_redistribute_regular_tensor(self, mock_device_mesh, mock_cp_mesh):
+        """Test redistribute_logits_for_cp with regular tensor input."""
+        from nemo_rl.models.automodel.train import redistribute_logits_for_cp
+
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+
+        logits = torch.randn(batch_size, seq_len, vocab_size)
+
+        # Mock device_mesh to return cp_mesh
+        mock_device_mesh.__getitem__ = MagicMock(return_value=mock_cp_mesh)
+
+        with patch(
+            "nemo_rl.models.automodel.train.DTensor.from_local"
+        ) as mock_from_local:
+            mock_dtensor = MagicMock()
+            mock_from_local.return_value = mock_dtensor
+
+            result = redistribute_logits_for_cp(
+                logits, mock_device_mesh, mock_cp_mesh, sequence_dim=1
+            )
+
+            # Verify DTensor.from_local was called with correct args
+            mock_from_local.assert_called_once()
+            call_args = mock_from_local.call_args
+            assert torch.equal(call_args[0][0], logits)
+
+    def test_redistribute_dtensor_input(self, mock_device_mesh, mock_cp_mesh):
+        """Test redistribute_logits_for_cp with DTensor input."""
+        from nemo_rl.models.automodel.train import redistribute_logits_for_cp
+
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+
+        # Create mock DTensor
+        mock_dtensor = MagicMock(spec=DTensor)
+        mock_dtensor.to_local.return_value = torch.randn(
+            batch_size, seq_len, vocab_size
+        )
+
+        # Mock device_mesh properties
+        mock_tp_mesh = MagicMock()
+        mock_tp_mesh.ndim = 1
+        mock_tp_mesh.mesh_dim_names = ["tp"]
+        mock_dtensor.device_mesh = mock_tp_mesh
+
+        mock_device_mesh.__getitem__ = MagicMock(return_value=mock_cp_mesh)
+
+        with patch(
+            "nemo_rl.models.automodel.train.DTensor.from_local"
+        ) as mock_from_local:
+            mock_result_dtensor = MagicMock()
+            mock_from_local.return_value = mock_result_dtensor
+
+            result = redistribute_logits_for_cp(
+                mock_dtensor, mock_device_mesh, mock_cp_mesh, sequence_dim=1
+            )
+
+            # Verify to_local was called
+            mock_dtensor.to_local.assert_called_once()
+
+            # Verify DTensor.from_local was called
+            mock_from_local.assert_called_once()
+
+
+# =====================
+# Test prepare_data_for_cp
+# =====================
+@pytest.mark.automodel
+class TestPrepareDataForCP:
+    def test_prepare_data_for_cp(self, mock_cp_mesh):
+        """Test prepare_data_for_cp function."""
+        from nemo_rl.models.automodel.train import prepare_data_for_cp
+
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+
+        # Create input data
+        input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+        position_ids = torch.arange(seq_len).repeat(batch_size, 1)
+        seq_index = torch.arange(seq_len).unsqueeze(0)
+
+        # Create processed inputs with cp_buffers
+        processed_inputs = ProcessedInputs(
+            input_ids=input_ids,
+            seq_len=seq_len,
+            attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+            position_ids=position_ids,
+            flash_attn_kwargs={},
+            vlm_kwargs={},
+            cp_buffers=[input_ids, position_ids],
+            seq_index=seq_index,
+        )
+
+        # Create data dict
+        mb = BatchedDataDict(
+            {
+                "input_ids": input_ids,
+                "position_ids": position_ids,
+                "other_tensor": torch.ones(batch_size, seq_len),
+            }
+        )
+
+        with patch(
+            "nemo_rl.models.automodel.train.DTensor.from_local"
+        ) as mock_from_local:
+            # Mock DTensor behavior
+            mock_dtensor = MagicMock()
+            mock_full_tensor = MagicMock()
+            mock_full_tensor.squeeze.return_value = seq_index.squeeze(0)
+            mock_dtensor.full_tensor.return_value = mock_full_tensor
+            mock_from_local.return_value = mock_dtensor
+
+            seq_index_result, updated_mb = prepare_data_for_cp(
+                mb, processed_inputs, mock_cp_mesh, sequence_dim=1
+            )
+
+            # Verify seq_index was added to mb
+            assert "seq_index" in updated_mb
+
+            # Verify DTensor.from_local was called for cp_buffers
+            assert mock_from_local.call_count >= 1
+
+
+# =====================
+# Test LogprobsPostProcessor with sequence packing
+# =====================
+@pytest.mark.automodel
+class TestLogprobsPostProcessorSeqPacking:
+    def test_logprobs_with_sequence_packing(
+        self, base_cfg, mock_device_mesh, mock_cp_mesh, mock_tp_mesh
+    ):
+        """Test LogprobsPostProcessor with sequence packing enabled."""
+        processor = LogprobsPostProcessor(
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+            enable_seq_packing=True,
+        )
+
+        original_batch_size = 4
+        original_seq_len = 32
+        packed_seq_len = 128  # All 4 sequences packed
+        vocab_size = 32000
+
+        logits = torch.randn(1, packed_seq_len, vocab_size)
+        input_ids = torch.randint(0, vocab_size, (1, packed_seq_len))
+        input_lengths = torch.tensor([32, 32, 32, 32])
+
+        @dataclass
+        class MockFlashAttnKwargs:
+            cu_seqlens_q: torch.Tensor
+
+        flash_kwargs = MockFlashAttnKwargs(
+            cu_seqlens_q=torch.tensor([0, 32, 64, 96, 128])
+        )
+
+        processed_inputs = ProcessedInputs(
+            input_ids=input_ids,
+            seq_len=packed_seq_len,
+            attention_mask=None,
+            position_ids=torch.arange(packed_seq_len).unsqueeze(0),
+            flash_attn_kwargs=flash_kwargs,
+            vlm_kwargs={},
+            cp_buffers=[],
+            seq_index=None,
+        )
+
+        result = processor(
+            logits=logits,
+            processed_inputs=processed_inputs,
+            input_lengths=input_lengths,
+            original_batch_size=original_batch_size,
+            original_seq_len=original_seq_len,
+            enable_seq_packing=True,
+        )
+
+        # Result should be unpacked to original shape
+        assert result.shape == (original_batch_size, original_seq_len)
+
+    def test_logprobs_masking_without_sequence_packing(
+        self, base_cfg, mock_device_mesh, mock_cp_mesh, mock_tp_mesh
+    ):
+        """Test LogprobsPostProcessor applies mask when sequence packing is disabled."""
+        processor = LogprobsPostProcessor(
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+            enable_seq_packing=False,
+        )
+
+        batch_size = 4
+        seq_len = 64
+        vocab_size = 32000
+
+        logits = torch.randn(batch_size, seq_len, vocab_size)
+        input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+        # Variable length sequences
+        input_lengths = torch.tensor([32, 48, 64, 16])
+
+        processed_inputs = ProcessedInputs(
+            input_ids=input_ids,
+            seq_len=seq_len,
+            attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+            position_ids=torch.arange(seq_len).repeat(batch_size, 1),
+            flash_attn_kwargs={},
+            vlm_kwargs={},
+            cp_buffers=[],
+            seq_index=None,
+        )
+
+        result = processor(
+            logits=logits,
+            processed_inputs=processed_inputs,
+            input_lengths=input_lengths,
+            original_batch_size=batch_size,
+            original_seq_len=seq_len,
+            enable_seq_packing=False,
+        )
+
+        # Verify result shape
+        assert result.shape == (batch_size, seq_len)
+
+        # Verify masking - positions beyond input_lengths should be zero
+        for i, length in enumerate(input_lengths):
+            # Positions beyond length should be zero
+            assert torch.all(result[i, length:] == 0)
+
+
+# =====================
+# Test TopkLogitsPostProcessor with sequence packing
+# =====================
+@pytest.mark.automodel
+class TestTopkLogitsPostProcessorSeqPacking:
+    def test_topk_with_sequence_packing(
+        self, base_cfg, mock_device_mesh, mock_cp_mesh, mock_tp_mesh
+    ):
+        """Test TopkLogitsPostProcessor with sequence packing enabled."""
+        k = 10
+        processor = TopkLogitsPostProcessor(
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+            k=k,
+            enable_seq_packing=True,
+        )
+
+        original_batch_size = 4
+        original_seq_len = 32
+        packed_seq_len = 128  # All 4 sequences packed
+        vocab_size = 32000
+
+        logits = torch.randn(1, packed_seq_len, vocab_size)
+        input_lengths = torch.tensor([32, 32, 32, 32])
+
+        @dataclass
+        class MockFlashAttnKwargs:
+            cu_seqlens_q: torch.Tensor
+
+        flash_kwargs = MockFlashAttnKwargs(
+            cu_seqlens_q=torch.tensor([0, 32, 64, 96, 128])
+        )
+
+        processed_inputs = ProcessedInputs(
+            input_ids=torch.randint(0, vocab_size, (1, packed_seq_len)),
+            seq_len=packed_seq_len,
+            attention_mask=None,
+            position_ids=torch.arange(packed_seq_len).unsqueeze(0),
+            flash_attn_kwargs=flash_kwargs,
+            vlm_kwargs={},
+            cp_buffers=[],
+            seq_index=None,
+        )
+
+        vals, idx = processor(
+            logits=logits,
+            processed_inputs=processed_inputs,
+            input_lengths=input_lengths,
+            original_batch_size=original_batch_size,
+            original_seq_len=original_seq_len,
+            enable_seq_packing=True,
+        )
+
+        # Result should be unpacked to original shape
+        assert vals.shape == (original_batch_size, original_seq_len, k)
+        assert idx.shape == (original_batch_size, original_seq_len, k)
+
+
+# =====================
+# Test automodel_forward_backward with backward pass
+# =====================
+@pytest.mark.automodel
+class TestAutomodelForwardBackwardWithGradients:
+    def test_forward_backward_computes_gradients(
+        self,
+        base_cfg,
+        mock_device_mesh,
+        mock_cp_mesh,
+        mock_tp_mesh,
+    ):
+        """Test automodel_forward_backward with forward_only=False computes gradients."""
+        batch_size = 2
+        seq_len = 16
+        vocab_size = 100
+        hidden_size = 32
+
+        # Create a simple model with trainable parameters
+        class SimpleModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.embed = torch.nn.Embedding(vocab_size, hidden_size)
+                self.proj = torch.nn.Linear(hidden_size, vocab_size)
+
+            def forward(self, input_ids, **kwargs):
+                x = self.embed(input_ids)
+                logits = self.proj(x)
+                return MagicMock(logits=logits)
+
+        model = SimpleModel()
+        model.train()
+
+        # Create processed inputs
+        input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+        processed_inputs = ProcessedInputs(
+            input_ids=input_ids,
+            seq_len=seq_len,
+            attention_mask=torch.ones(batch_size, seq_len, dtype=torch.bool),
+            position_ids=torch.arange(seq_len).repeat(batch_size, 1),
+            flash_attn_kwargs={},
+            vlm_kwargs={},
+            cp_buffers=[],
+            seq_index=None,
+        )
+
+        # Create data dict
+        data_dict = BatchedDataDict(
+            {
+                "input_ids": input_ids,
+                "input_lengths": torch.full((batch_size,), seq_len),
+                "sample_mask": torch.ones(batch_size, dtype=torch.bool),
+            }
+        )
+
+        # Create processed microbatch
+        processed_mb = ProcessedMicrobatch(
+            data_dict=data_dict,
+            processed_inputs=processed_inputs,
+            original_batch_size=batch_size,
+            original_seq_len=seq_len,
+        )
+
+        # Create loss function that returns requires_grad tensor
+        def loss_fn(logits, mb, global_valid_seqs, global_valid_toks):
+            loss = logits.mean()
+            return loss, {"loss": loss.item()}
+
+        # Create loss post-processor
+        loss_post_processor = LossPostProcessor(
+            loss_fn=loss_fn,
+            cfg=base_cfg,
+            device_mesh=mock_device_mesh,
+            cp_mesh=mock_cp_mesh,
+            tp_mesh=mock_tp_mesh,
+            cp_size=1,
+            dp_size=1,
+        )
+
+        # Verify no gradients initially
+        assert model.proj.weight.grad is None
+
+        # Call automodel_forward_backward with forward_only=False
+        results = automodel_forward_backward(
+            model=model,
+            cfg=base_cfg,
+            data_iterator=iter([processed_mb]),
+            post_processing_fn=loss_post_processor,
+            forward_only=False,  # Enable backward pass
+            global_valid_seqs=torch.tensor(batch_size),
+            global_valid_toks=torch.tensor(batch_size * seq_len),
+            dp_size=1,
+            cp_size=1,
+        )
+
+        # Verify gradients were computed
+        assert model.proj.weight.grad is not None
+        assert len(results) == 1
+
+
+# =====================
+# Test aggregate_training_statistics
+# =====================
+@pytest.mark.automodel
+class TestAggregateTrainingStatistics:
+    """Tests for aggregate_training_statistics function."""
+
+    @patch("torch.distributed.all_reduce")
+    @patch("torch.distributed.get_rank", return_value=0)
+    @patch("torch.cuda.get_device_name", return_value="NVIDIA A100")
+    def test_basic_aggregation(self, mock_device_name, mock_get_rank, mock_all_reduce):
+        """Test basic statistics aggregation."""
+        from nemo_rl.models.automodel.train import (
+            aggregate_training_statistics,
+        )
+
+        losses = [0.5, 0.4, 0.3]
+        all_mb_metrics = [
+            {"loss": 0.5, "num_valid_samples": 4},
+            {"loss": 0.4, "num_valid_samples": 4},
+            {"loss": 0.3, "num_valid_samples": 4},
+        ]
+        grad_norm = torch.tensor([1.5])
+        mock_dp_group = MagicMock()
+        dtype = torch.bfloat16
+
+        metrics = aggregate_training_statistics(
+            losses=losses,
+            all_mb_metrics=all_mb_metrics,
+            grad_norm=grad_norm,
+            dp_group=mock_dp_group,
+            dtype=dtype,
+        )
+
+        # Verify all_reduce was called
+        mock_all_reduce.assert_called_once()
+
+        # Verify metrics structure
+        assert "global_loss" in metrics
+        assert "grad_norm" in metrics
+        assert "rank" in metrics
+        assert "gpu_name" in metrics
+        assert "model_dtype" in metrics
+        assert "all_mb_metrics" in metrics
+
+        # Verify values
+        assert metrics["rank"] == 0
+        assert metrics["gpu_name"] == "NVIDIA A100"
+        assert metrics["model_dtype"] == torch.bfloat16
+        assert torch.equal(metrics["grad_norm"], grad_norm)
+
+        # Verify metrics aggregation
+        assert metrics["all_mb_metrics"]["loss"] == [0.5, 0.4, 0.3]
+        assert metrics["all_mb_metrics"]["num_valid_samples"] == [4, 4, 4]
+
+    @patch("torch.distributed.all_reduce")
+    @patch("torch.distributed.get_rank", return_value=2)
+    @patch("torch.cuda.get_device_name", return_value="NVIDIA H100")
+    def test_with_none_grad_norm(
+        self, mock_device_name, mock_get_rank, mock_all_reduce
+    ):
+        """Test aggregation with None grad_norm (eval mode)."""
+        from nemo_rl.models.automodel.train import (
+            aggregate_training_statistics,
+        )
+
+        losses = [0.5]
+        all_mb_metrics = [{"loss": 0.5}]
+        grad_norm = None
+        mock_dp_group = MagicMock()
+        dtype = torch.float16
+
+        metrics = aggregate_training_statistics(
+            losses=losses,
+            all_mb_metrics=all_mb_metrics,
+            grad_norm=grad_norm,
+            dp_group=mock_dp_group,
+            dtype=dtype,
+        )
+
+        assert metrics["grad_norm"] is None
+        assert metrics["rank"] == 2
+        assert metrics["model_dtype"] == torch.float16
+
+    @patch("torch.distributed.all_reduce")
+    @patch("torch.distributed.get_rank", return_value=0)
+    @patch("torch.cuda.get_device_name", return_value="NVIDIA A100")
+    def test_empty_metrics(self, mock_device_name, mock_get_rank, mock_all_reduce):
+        """Test aggregation with empty metrics list."""
+        from nemo_rl.models.automodel.train import (
+            aggregate_training_statistics,
+        )
+
+        losses = []
+        all_mb_metrics = []
+        grad_norm = torch.tensor([0.0])
+        mock_dp_group = MagicMock()
+        dtype = torch.bfloat16
+
+        metrics = aggregate_training_statistics(
+            losses=losses,
+            all_mb_metrics=all_mb_metrics,
+            grad_norm=grad_norm,
+            dp_group=mock_dp_group,
+            dtype=dtype,
+        )
+
+        assert metrics["all_mb_metrics"] == {}
+        assert metrics["global_loss"].numel() == 0
+
+    @patch("torch.distributed.all_reduce")
+    @patch("torch.distributed.get_rank", return_value=0)
+    @patch("torch.cuda.get_device_name", return_value="NVIDIA A100")
+    def test_global_loss_on_cpu(self, mock_device_name, mock_get_rank, mock_all_reduce):
+        """Test that global_loss is moved to CPU."""
+        from nemo_rl.models.automodel.train import (
+            aggregate_training_statistics,
+        )
+
+        losses = [0.5, 0.4]
+        all_mb_metrics = [{"loss": 0.5}, {"loss": 0.4}]
+        grad_norm = torch.tensor([1.0])
+        mock_dp_group = MagicMock()
+        dtype = torch.bfloat16
+
+        metrics = aggregate_training_statistics(
+            losses=losses,
+            all_mb_metrics=all_mb_metrics,
+            grad_norm=grad_norm,
+            dp_group=mock_dp_group,
+            dtype=dtype,
+        )
+
+        # Verify global_loss is on CPU
+        assert metrics["global_loss"].device == torch.device("cpu")