Train batch generic (#724)

Co-authored-by: Hossein Kavianihamedani <hosseinkh@fb.com>
Co-authored-by: Felipe Mello <fmellomascarenhas@gmail.com>

Author: 3 people

apps/grpo/main.py

@@ -332,8 +332,7 @@ async def continuous_training():
             else:
                 t.step("waiting_for_buffer")
 
-                inputs, targets = batch
-                await trainer.train_step.call(inputs, targets)
+                await trainer.train_step.call(batch)
                 training_step += 1
                 t.step("train_step")
 

src/forge/actors/trainer/titan.py

@@ -20,6 +20,7 @@
 from forge.observability.metrics import record_metric, Reduce
 from forge.observability.perf_tracker import Tracer
 from forge.rl.loss import create_shifted_targets
+from forge.types import TrainBatch
 from monarch.actor import endpoint
 from torch import Tensor
 from torch.distributed.checkpoint._nested_dict import flatten_state_dict
@@ -117,17 +118,15 @@ async def setup(self):
         self.engine.checkpointer.load(step=self.step)
         self.engine.optimizers.zero_grad()
 
-    def forward_backward(
-        self, inputs: dict[str, Tensor], targets: dict[str, Tensor]
-    ) -> Tensor:
+    def forward_backward(self, batch: TrainBatch) -> Tensor:
         model_parts = self.engine.model_parts
         parallel_dims = self.engine.parallel_dims
         optional_context_parallel_ctx = None
 
         # Create shifted target_ids for next-token prediction
         # target_ids[i] = input_ids[i+1], with loss_mask applied
-        targets["target_ids"] = create_shifted_targets(
-            inputs["tokens"], targets.get("loss_mask")
+        batch.loss_inputs["target_ids"] = create_shifted_targets(
+            batch.model_inputs["tokens"], batch.loss_inputs.get("loss_mask")
         )
 
         if parallel_dims.pp_enabled:
@@ -136,8 +135,8 @@ def forward_backward(
             with self.engine.train_context(optional_context_parallel_ctx):
                 assert len(model_parts) == 1
                 with self.engine.maybe_enable_amp:
-                    logits = model_parts[0](**inputs)
-                    loss_output = self.loss(logits, **targets)
+                    logits = model_parts[0](**batch.model_inputs)
+                    loss_output = self.loss(logits, **batch.loss_inputs)
                     loss = loss_output.loss
 
                 # Record metrics from loss output
@@ -156,19 +155,16 @@ def forward_backward(
         return loss
 
     @endpoint
-    async def train_step(
-        self, inputs: list[dict[str, Tensor]], targets: list[dict[str, Tensor]]
-    ) -> float:
+    async def train_step(self, batches: list[TrainBatch]) -> float:
         t = Tracer("rl_trainer_perf/step", timer="gpu", track_memory=True)
         t.start()
 
         self.engine.gc_handler.run(self.step)
-        local_inputs = inputs[self.engine.dp_rank]
-        local_targets = targets[self.engine.dp_rank]
-        batch_to_device(local_inputs, self.engine.device)
-        batch_to_device(local_targets, self.engine.device)
+        batch = batches[self.engine.dp_rank]
+        batch_to_device(batch.model_inputs, self.engine.device)
+        batch_to_device(batch.loss_inputs, self.engine.device)
 
-        loss = self.forward_backward(local_inputs, local_targets)
+        loss = self.forward_backward(batch)
         torch.distributed.all_reduce(loss)
 
         t.step("forward_backward")

src/forge/rl/collate.py

@@ -4,21 +4,17 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Any
-
 import torch
 from forge.rl.types import Group
+from forge.types import TrainBatch
 
 
-def collate(
-    batches: list[Group],
-) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
+def collate(batches: list[Group]) -> list[TrainBatch]:
     """
-    Collates a list of batches into a single batch of inputs and targets.
+    Collates a list of batches into TrainBatch objects.
     Each batch is a list of episodes, and each episode is a dict of tensors.
     """
-    inputs = []
-    targets = []
+    result = []
     for batch in batches:
         request = [e.request_tensor for e in batch]
         request = torch.stack(request)  # [b x s]
@@ -41,14 +37,18 @@ def collate(
         generator_logprobs = torch.stack([e.generator_logprobs for e in batch])
         loss_mask = torch.stack([e.loss_mask for e in batch])
 
-        input = {"tokens": input_ids}
-        target = {
+        loss_inputs = {
             "generator_logprobs": generator_logprobs,
             "loss_mask": loss_mask,
             "advantages": advantages,
         }
         if ref_logprobs is not None:
-            target["ref_logprobs"] = ref_logprobs
-        inputs.append(input)
-        targets.append(target)
-    return inputs, targets
+            loss_inputs["ref_logprobs"] = ref_logprobs
+
+        result.append(
+            TrainBatch(
+                model_inputs={"tokens": input_ids},
+                loss_inputs=loss_inputs,
+            )
+        )
+    return result

src/forge/types.py

@@ -126,3 +126,34 @@ class ProvisionerConfig:
     """A config for the forge provisioner."""
 
     launcher_config: LauncherConfig
+
+
+@dataclass
+class TrainBatch:
+    """Universal training batch for all Forge training modes.
+
+    Usage:
+        logits = model(**batch.model_inputs)
+        loss = loss_fn(logits, **batch.loss_inputs)
+
+    Attributes:
+        model_inputs (dict[str, Any]): Inputs for model forward pass (e.g., input_ids, attention_mask).
+        loss_inputs (dict[str, Any]): Inputs for loss computation (e.g., target_ids, advantages, beta).
+        meta (dict[str, Any]): Any extra metadata that is not a model or loss input.
+
+    Example:
+        >>> # SFT
+        >>> batch = TrainBatch(
+        >>>     model_inputs={"input_ids": ids, "attention_mask": mask},
+        >>>     loss_inputs={"target_ids": targets},
+        >>> )
+        >>> # RL (GRPO)
+        >>> batch = TrainBatch(
+        >>>     model_inputs={"input_ids": ids},
+        >>>     loss_inputs={"target_ids": targets, "advantages": adv, "ref_logprobs": ref},
+        >>> )
+    """
+
+    model_inputs: dict[str, Any]
+    loss_inputs: dict[str, Any]
+    meta: dict[str, Any] = field(default_factory=dict)

tests/sandbox/rl_trainer/main.py

@@ -21,6 +21,7 @@
     ProcessConfig,
     ProvisionerConfig,
     ServiceConfig,
+    TrainBatch,
 )
 from forge.util.config import parse
 from omegaconf import DictConfig
@@ -75,13 +76,12 @@ def generate_random_batch(
     vocab_size: int = 32000,
     device: str = "cuda",
     dp_size: int = 1,
-):
+) -> list[TrainBatch]:
     """
-    Generate random input and target tensors matching GRPO data format
-    Creates one batch per data parallel rank
+    Generate random TrainBatch objects matching GRPO data format.
+    Creates one batch per data parallel rank.
     """
-    inputs = []
-    targets = []
+    batches = []
 
     # Create one batch for each data parallel rank
     for _ in range(dp_size):
@@ -109,17 +109,19 @@ def generate_random_batch(
         )
         advantages = torch.randn((local_batch_size, 1), device=device)
         input_tokens = torch.cat([request, response], dim=1)
-        inputs.append({"tokens": input_tokens})
-        targets.append(
-            {
-                "response": response,
-                "ref_logprobs": ref_logprobs,
-                "advantages": advantages,
-                "padding_mask": padding_mask,
-            }
+        batches.append(
+            TrainBatch(
+                model_inputs={"tokens": input_tokens},
+                loss_inputs={
+                    "response": response,
+                    "ref_logprobs": ref_logprobs,
+                    "advantages": advantages,
+                    "padding_mask": padding_mask,
+                },
+            )
         )
 
-    return inputs, targets
+    return batches
 
 
 async def main(cfg: DictConfig):
@@ -201,7 +203,7 @@ async def continuous_training():
             t = Tracer("trainer/continuous_training")
             t.start()
 
-            inputs, targets = generate_random_batch(
+            batches = generate_random_batch(
                 local_batch_size=local_batch_size,
                 request_len=request_len,
                 response_len=response_len,
@@ -211,7 +213,7 @@ async def continuous_training():
             t.step("generate_random_data")
 
             # Perform training step
-            await trainer.train_step.call(inputs, targets)
+            await trainer.train_step.call(batches)
             training_step += 1
             t.step("train_step")
 

tests/sandbox/weight_sync/main.py

@@ -24,7 +24,7 @@
 from forge.actors.trainer import RLTrainer
 from forge.controller.provisioner import init_provisioner, shutdown
 from forge.observability.metric_actors import get_or_create_metric_logger
-from forge.types import LauncherConfig, ProvisionerConfig
+from forge.types import LauncherConfig, ProvisionerConfig, TrainBatch
 from forge.util.config import parse
 from omegaconf import DictConfig
 from vllm.transformers_utils.tokenizer import get_tokenizer
@@ -37,13 +37,12 @@ def generate_random_batch(
     vocab_size: int = 32000,
     device: str = "cuda",
     dp_size: int = 1,
-):
+) -> list[TrainBatch]:
     """
-    Generate random input and target tensors for a single training step.
+    Generate random TrainBatch objects for a single training step.
     Creates one batch per data parallel rank.
     """
-    inputs = []
-    targets = []
+    batches = []
 
     # Create one batch for each data parallel rank
     for _ in range(dp_size):
@@ -71,17 +70,19 @@ def generate_random_batch(
         )
         advantages = torch.randn((local_batch_size, 1), device=device)
         input_tokens = torch.cat([request, response], dim=1)
-        inputs.append({"tokens": input_tokens})
-        targets.append(
-            {
-                "response": response,
-                "ref_logprobs": ref_logprobs,
-                "advantages": advantages,
-                "padding_mask": padding_mask,
-            }
+        batches.append(
+            TrainBatch(
+                model_inputs={"tokens": input_tokens},
+                loss_inputs={
+                    "response": response,
+                    "ref_logprobs": ref_logprobs,
+                    "advantages": advantages,
+                    "padding_mask": padding_mask,
+                },
+            )
         )
 
-    return inputs, targets
+    return batches
 
 
 async def main(cfg: DictConfig):
@@ -147,15 +148,15 @@ async def main(cfg: DictConfig):
     print("Running single training step...")
     step_start = time.time()
 
-    inputs, targets = generate_random_batch(
+    batches = generate_random_batch(
         local_batch_size=local_batch_size,
         request_len=request_len,
         response_len=response_len,
         vocab_size=vocab_size,
         dp_size=dp_size,
     )
 
-    await trainer.train_step.call(inputs, targets)
+    await trainer.train_step.call(batches)
     step_time = time.time() - step_start
     print(f"Finished train step in ({step_time:.2f}s)\n")