Add Dr. GRPO loss normalization

src/art/_backend_training.py

@@ -34,9 +34,27 @@
 _GRADIENT_TRAIN_TIME = "time/gradient_step_train_s"
 
 
+def _validate_loss_reduction(
+    loss_type: dev.LossType | None,
+    max_completion_length: int | None,
+) -> None:
+    if loss_type is not None and (
+        not isinstance(loss_type, str) or loss_type not in {"grpo", "bnpo", "dr_grpo"}
+    ):
+        raise ValueError("loss_type must be one of 'grpo', 'bnpo', or 'dr_grpo'.")
+    if max_completion_length is not None and (
+        isinstance(max_completion_length, bool)
+        or not isinstance(max_completion_length, int)
+        or max_completion_length < 1
+    ):
+        raise ValueError("max_completion_length must be a positive integer")
+
+
 def build_rl_train_configs(
     *,
     learning_rate: float,
+    loss_type: dev.LossType | None = None,
+    max_completion_length: int | None = None,
     advantage_balance: float = 0.0,
     scale_rewards: bool = True,
     importance_sampling_level: Literal[
@@ -63,6 +81,9 @@ def build_rl_train_configs(
     final_training_step: int | None = None,
     grad_accumulation_sequences: int | None = None,
 ) -> tuple[TrainConfig, dev.TrainConfig]:
+    if loss_type == "dr_grpo" and max_completion_length is None:
+        max_completion_length = dev.DEFAULT_MAX_COMPLETION_LENGTH
+    _validate_loss_reduction(loss_type, max_completion_length)
     config = TrainConfig(
         learning_rate=learning_rate,
         kl_penalty_coef=kl_penalty_coef,
@@ -82,6 +103,14 @@ def build_rl_train_configs(
         "scale_rewards": scale_rewards,
     }
 
+    # Keep these optional so existing callers retain ART's historical BNPO
+    # behavior, while allowing a backend-level setting to override the
+    # trainer's native GRPOConfig when explicitly requested.
+    if loss_type is not None:
+        dev_config["loss_type"] = loss_type
+    if max_completion_length is not None:
+        dev_config["max_completion_length"] = max_completion_length
+
     if allow_training_without_logprobs is not None:
         dev_config["allow_training_without_logprobs"] = allow_training_without_logprobs
     if plot_tensors is not None:

src/art/dev/__init__.py

@@ -13,7 +13,12 @@
     VllmRuntimeArgs,
 )
 from .openai_server import OpenAIServerConfig, ServerArgs, get_openai_server_config
-from .train import TrainConfig, TrainSFTConfig
+from .train import (
+    DEFAULT_MAX_COMPLETION_LENGTH,
+    LossType,
+    TrainConfig,
+    TrainSFTConfig,
+)
 from .validate import (
     is_dedicated_mode,
     is_external_vllm_mode,
@@ -40,5 +45,7 @@
     "ServerArgs",
     "TrainSFTConfig",
     "TrainConfig",
+    "LossType",
+    "DEFAULT_MAX_COMPLETION_LENGTH",
     "validate_dedicated_config",
 ]

src/art/dev/model.py

@@ -4,6 +4,7 @@
 from typing_extensions import Required, TypedDict
 
 from .engine import EngineArgs
+from .train import LossType
 
 RolloutWeightUpdateMode = Literal["step_lora", "in_flight_lora"]
 VllmRuntimeMode = Literal["managed", "external"]
@@ -378,6 +379,7 @@ class TrainerArgs(TypedDict, total=False):
     num_generations: int | None
     temperature: float
     max_completion_length: int | None
+    loss_type: LossType
     ds3_gather_for_generation: bool
     beta: float
     reward_weights: list[float] | None

src/art/dev/train.py

@@ -8,7 +8,21 @@
     from art.megatron.routing_replay import MoeRoutingReplayBundle
 
 
+# Reduction schemes implemented by the RL policy-loss reducer.  Keep this
+# alias in the developer configuration module so the public ``TrainerArgs``
+# and the backend's experimental configuration cannot drift apart.
+LossType = Literal["grpo", "bnpo", "dr_grpo"]
+
+# TRL 0.20 (the backend version pinned by ART) uses this value when callers
+# select ``loss_type="dr_grpo"`` without spelling out a completion length.
+# Keeping the default here lets the public TrainerArgs API behave exactly like
+# the native GRPOConfig while still allowing an explicit override.
+DEFAULT_MAX_COMPLETION_LENGTH = 256
+
+
 class TrainConfig(TypedDict, total=False):
+    loss_type: LossType
+    max_completion_length: int
     advantage_balance: float
     """Balance between negative and positive advantages in the range [-1.0, 1.0]. \
 -1.0 means only training on negative advantages, 1.0 means only training on \

src/art/local/backend.py

@@ -1291,6 +1291,8 @@ async def train(  # type: ignore[override]
         # Core training parameters
         learning_rate: float = 5e-6,
         loss_fn: Literal["cispo", "ppo"] = "cispo",
+        loss_type: dev.LossType | None = None,
+        max_completion_length: int | None = None,
         loss_fn_config: dict | None = None,
         normalize_advantages: bool = True,
         adam_params: object | None = None,
@@ -1341,6 +1343,14 @@ async def train(  # type: ignore[override]
             learning_rate: Learning rate for training. Defaults to 5e-6.
             loss_fn: RL loss function. LocalBackend currently supports
                 "cispo" and "ppo".
+            loss_type: Token-loss reduction. Use ``"grpo"`` for per-completion
+                means, ``"bnpo"`` for ART's historical active-token mean, or
+                ``"dr_grpo"`` for the fixed response-length denominator from
+                the Dr. GRPO paper. If omitted, the Unsloth trainer's
+                ``TrainerArgs.loss_type`` is used when available.
+            max_completion_length: Fixed completion length used by
+                ``loss_type="dr_grpo"``. Defaults to TRL's 256-token value
+                when that reduction is selected without an explicit length.
             loss_fn_config: Additional loss-function config. Not supported by
                 LocalBackend.
             normalize_advantages: Backward-compatible alias for reward std scaling.
@@ -1404,6 +1414,15 @@ async def train(  # type: ignore[override]
             # await model.log(metrics=result.metrics, step=result.step)
         """
         groups_list = list(trajectory_groups)
+        configured_trainer_args = (
+            (model._internal_config or {}).get("trainer_args", {})
+            if isinstance(model, TrainableModel)
+            else {}
+        )
+        if loss_type is None:
+            loss_type = configured_trainer_args.get("loss_type")
+        if max_completion_length is None:
+            max_completion_length = configured_trainer_args.get("max_completion_length")
         if loss_fn not in {"cispo", "ppo"}:
             raise ValueError("LocalBackend only supports loss_fn='cispo' or 'ppo'.")
         if loss_fn_config is not None:
@@ -1441,6 +1460,8 @@ async def train(  # type: ignore[override]
             )
         config, dev_config = build_rl_train_configs(
             learning_rate=learning_rate,
+            loss_type=loss_type,
+            max_completion_length=max_completion_length,
             advantage_balance=advantage_balance,
             scale_rewards=scale_rewards,
             importance_sampling_level=importance_sampling_level,