PullRequest: 998 新增Off-Policy配置文件并调整相关训练参数

Merge branch chucai.dzq/align-offpolicy of [email protected]:inclusionAI/AReaL.git into asystem/gh
https://code.alipay.com/inclusionAI/AReaL/pull_requests/998

Reviewed-by: 峯回 <[email protected]>


* align off policy

Author: 楚财

areal/api/alloc_mode.py

@@ -249,19 +249,20 @@ async def _async_custom_function_call(self, method: str, *args, **kwargs):
 
     def _dispatch_inputs(self, *args, **kwargs):
         # Find and split DistributedBatch arguments
+        rebalance = kwargs.pop("rebalance", True)
         split_args = []
         for arg in args:
             if isinstance(arg, DistributedBatch):
                 # Split across DP groups
-                split_args.append(self._align_batches_with_dp(arg, rebalance=True))
+                split_args.append(self._align_batches_with_dp(arg, rebalance=rebalance))
             else:
                 # Replicate to all DP heads
                 split_args.append([arg] * self.parallel_strategy.dp_size)
 
         split_kwargs = {}
         for k, v in kwargs.items():
             if isinstance(v, DistributedBatch):
-                split_kwargs[k] = self._align_batches_with_dp(v, rebalance=True)
+                split_kwargs[k] = self._align_batches_with_dp(v, rebalance=rebalance)
             else:
                 split_kwargs[k] = [v] * self.parallel_strategy.dp_size
         return split_args, split_kwargs
@@ -633,7 +634,7 @@ def train_batch(
             gradient norm, etc.
         """
         return self._custom_function_call(
-            "train_batch", input_, loss_fn, loss_weight_fn
+            "train_batch", input_, loss_fn, loss_weight_fn, rebalance=True
         )
 
     def eval_batch(