[feat] add --critic-lr and --num-critic-only-steps

Author: zhuzilin

slime/backends/megatron_utils/actor.py

@@ -53,6 +53,7 @@ def init(self, args, role, wandb_run_id, with_ref=False):
         if role == "critic":
             self.args.load = self.args.critic_load
             self.args.save = self.args.critic_save
+            self.args.lr = self.args.critic_lr
 
         (self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id) = initialize_model_and_optimizer(
             args, role
@@ -254,11 +255,16 @@ def train_critic(self, rollout_id, rollout_data):
             self.model,
             data_iterator,
             num_microbatches,
-        )
-        values = [value.squeeze(-1) for value in values["values"]]
-        values, log_probs, ref_log_probs = sync_actor_critic_data(
-            self.args, values, None, None, self._actor_critic_groups
-        )
+        )["values"]
+
+        if rollout_id < self.args.num_critic_only_steps:
+            # we will only use the shape of log_probs in this situation
+            log_probs = values
+            ref_log_probs = values
+        else:
+            values, log_probs, ref_log_probs = sync_actor_critic_data(
+                self.args, values, None, None, self._actor_critic_groups
+            )
 
         rollout_data.update(
             {

slime/backends/megatron_utils/data.py

@@ -5,10 +5,10 @@
 import torch
 import torch.distributed as dist
 import torch.nn.functional as F
-import wandb
 from megatron.core import mpu
 from megatron.core.packed_seq_params import PackedSeqParams
 
+import wandb
 from slime.utils.data import get_minimum_num_micro_batch_size
 from slime.utils.flops_utils import calculate_fwd_flops
 from slime.utils.seqlen_balancing import get_seqlen_balanced_partitions
@@ -243,7 +243,7 @@ def log_rollout_data(rollout_id, args, rollout_data):
                     # NOTE: Here we have to do the clone().detach(), otherwise the tensor will be
                     # modified in place and will cause problem for the next rollout.
                     val = torch.cat(val).clone().detach()
-                    if key in ["log_probs", "ref_log_probs", "rollout_log_probs", "returns", "advantages"]:
+                    if key in ["log_probs", "ref_log_probs", "rollout_log_probs", "returns", "advantages", "values"]:
                         sum_of_sample_mean = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks)
                         val = cp_size * sum_of_sample_mean(val) / len(loss_masks)
                     else:

slime/backends/megatron_utils/loss.py

@@ -117,7 +117,7 @@ def get_values(
         response_lengths=response_lengths,
     ):
         assert logits_chunk.size(-1) == 1, f"{logits_chunk.shape}"
-        value_list.append(logits_chunk)
+        value_list.append(logits_chunk.squeeze(-1))
 
     return {
         "values": value_list,
@@ -366,19 +366,22 @@ def value_loss_function(args, batch, logits, sum_of_sample_mean):
 
     returns = torch.cat(batch["returns"], dim=0)
 
+    values_clipfrac = torch.abs(values - old_values) > args.value_clip
     values_clipped = old_values + (values - old_values).clamp(-args.value_clip, args.value_clip)
     surr1 = (values_clipped - returns) ** 2
     surr2 = (values - returns) ** 2
     loss = torch.max(surr1, surr2)
 
     loss = sum_of_sample_mean(loss)
+    values_clipfrac = sum_of_sample_mean(values_clipfrac.float())
 
     # make sure the gradient could backprop correctly.
     if values.numel() == 0:
         loss += 0 * values.sum()
 
     reported_loss = {
         "value_loss": loss.clone().detach(),
+        "value_clipfrac": values_clipfrac.clone().detach(),
     }
 
     return loss, reported_loss

slime/backends/megatron_utils/model.py

@@ -462,8 +462,10 @@ def train(rollout_id, model, optimizer, opt_param_scheduler, data_iterator, num_
                 for key, val in loss_dict.items()
             }
             log_dict["train/grad_norm"] = grad_norm
+            role = getattr(model[0], "role", "actor")
+            role_tag = "" if role == "actor" else f"{role}-"
             for param_group_id, param_group in enumerate(optimizer.param_groups):
-                log_dict[f"train/lr-pg_{param_group_id}"] = opt_param_scheduler.get_lr(param_group)
+                log_dict[f"train/{role_tag}lr-pg_{param_group_id}"] = opt_param_scheduler.get_lr(param_group)
 
             if args.use_wandb:
                 log_dict["train/step"] = accumulated_step_id
@@ -475,7 +477,7 @@ def train(rollout_id, model, optimizer, opt_param_scheduler, data_iterator, num_
                 if accumulated_step_id == 0 and "train/kl_loss" in log_dict:
                     assert log_dict["train/kl_loss"] == 0.0
 
-            print(f"step {accumulated_step_id}: {log_dict}")
+            print(f"{role_tag}step {accumulated_step_id}: {log_dict}")
     # Close out pre-hooks if using distributed optimizer and overlapped param gather.
     if pre_hook_enabled:
         disable_forward_pre_hook(model)
@@ -501,6 +503,7 @@ def save(iteration, model, optimizer, opt_param_scheduler):
 
 def initialize_model_and_optimizer(args, role: str = "actor"):
     model, optimizer, opt_param_scheduler = setup_model_and_optimizer(args, role)
+    setattr(model[0], "role", role)
     clear_memory()
     iteration, _ = load_checkpoint(
         model,

slime/backends/megatron_utils/model_provider.py

@@ -30,6 +30,10 @@ def __init__(
         if self.sequence_parallel:
             self.weight.sequence_parallel = True
 
+        self.weight.data.normal_(mean=0.0, std=0.02)
+        if bias:
+            self.bias.data.zero_()
+
     def forward(
         self,
         input_,

slime/utils/arguments.py

@@ -518,9 +518,12 @@ def add_algo_arguments(parser):
             reset_arg(parser, "--seed", type=int, default=1234)
             reset_arg(parser, "--clip-grad", type=float, default=1.0)
             reset_arg(parser, "--calculate-per-token-loss", action="store_true")
+            reset_arg(parser, "--lr", type=float, default=1e-6)
 
+            parser.add_argument("--num-critic-only-steps", type=int, default=0, help="Number of critic only steps")
             parser.add_argument("--critic-load", type=str, default=None, help="The checkpoint for critic model.")
             parser.add_argument("--critic-save", type=str, default=None, help="The checkpoint for critic model.")
+            parser.add_argument("--critic-lr", type=float, default=None, help="The lr for critic model")
 
             parser.add_argument("--eps-clip", type=float, default=0.2, help="PPO clip range")
             parser.add_argument("--eps-clip-high", type=float, default=None, help="PPO clip upper range")
@@ -984,9 +987,6 @@ def slime_validate_args(args):
             args.ckpt_step = args.ref_ckpt_step
         args.start_rollout_id = 0
 
-    if args.critic_load is None:
-        args.critic_load = args.load
-
     if args.eval_interval is not None:
         assert args.eval_prompt_data is not None, "eval_prompt_data must be set when eval_interval is set"
         if len(args.eval_prompt_data) == 1:
@@ -1032,6 +1032,10 @@ def slime_validate_args(args):
         args.critic_num_gpus_per_node = args.actor_num_gpus_per_node
     if args.critic_num_nodes is None:
         args.critic_num_nodes = args.actor_num_nodes
+    if args.critic_load is None:
+        args.critic_load = args.load
+    if args.critic_lr is None:
+        args.critic_lr = args.lr
 
     if args.debug_rollout_only:
         if args.colocate and args.rollout_num_gpus is None:

slime/utils/http_utils.py

@@ -83,7 +83,7 @@ def init_http_client(concurrency: int):
     if _http_client is None:
         _http_client = httpx.AsyncClient(
             limits=httpx.Limits(max_connections=concurrency),
-            timeout=httpx.Timeout(None),
+            timeout=httpx.Timeout(None, connect=5.0),
         )
 
 
@@ -113,7 +113,6 @@ async def post(url, payload, max_retries=60):
 
 
 async def get(url):
-    # never timeout
     response = await _http_client.get(url)
     response.raise_for_status()
     output = response.json()

train.py

@@ -70,11 +70,11 @@ def train(args):
 
         if args.use_critic:
             critic_train_handle = critic_model.async_train(rollout_id, rollout_data_ref)
-
-        ray.get(actor_model.async_train(rollout_id, rollout_data_ref))
-
-        if args.use_critic:
+            if rollout_id >= args.num_critic_only_steps:
+                ray.get(actor_model.async_train(rollout_id, rollout_data_ref))
             ray.get(critic_train_handle)
+        else:
+            ray.get(actor_model.async_train(rollout_id, rollout_data_ref))
 
         if args.save_interval is not None and (
             (rollout_id + 1) % args.save_interval == 0
@@ -85,9 +85,12 @@ def train(args):
                 ray.get(rollout_manager.save.remote(rollout_id))
 
         if args.offload:
-            ray.get(actor_model.async_offload())
             if args.use_critic:
                 ray.get(critic_model.async_offload())
+                if rollout_id >= args.num_critic_only_steps:
+                    ray.get(actor_model.async_offload())
+            else:
+                ray.get(actor_model.async_offload())
 
             ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS]))
 

train_async.py

@@ -28,10 +28,13 @@ def train(args):
     assert args.num_rollout > 0
 
     # sync the initialization (model initalization, load checkpoint, etc.)
-    # Note that we initialize it earlier as megatron ckpt loading may have really large peak memory usage.
+    if args.use_critic:
+        critic_init_handle = critic_model.async_init(args, role="critic", with_ref=False)
+
     start_rollout_ids = ray.get(
         actor_model.async_init(args, role="actor", with_ref=args.kl_coef != 0 or args.use_kl_loss)
     )
+
     assert len(set(start_rollout_ids)) == 1
     if args.start_rollout_id is None:
         args.start_rollout_id = start_rollout_ids[0]
@@ -42,6 +45,10 @@ def train(args):
     # initialize the connection for weight update during training
     ray.get(actor_model.async_init_weight_update_connections(rollout_manager))
 
+    if args.use_critic:
+        ray.get(critic_init_handle)
+        ray.get(actor_model.async_connect(critic_model))
+
     # always update weight first so that sglang has the loaded weights from training.
     ray.get(actor_model.async_update_weights())
 
@@ -56,7 +63,13 @@ def train(args):
         if rollout_id + 1 < args.num_rollout:
             rollout_data_next_future = rollout_manager.generate.remote(rollout_id + 1)
 
-        ray.get(actor_model.async_train(rollout_id, rollout_data_curr_ref))
+        if args.use_critic:
+            critic_train_handle = critic_model.async_train(rollout_id, rollout_data_curr_ref)
+            if rollout_id >= args.num_critic_only_steps:
+                ray.get(actor_model.async_train(rollout_id, rollout_data_curr_ref))
+            ray.get(critic_train_handle)
+        else:
+            ray.get(actor_model.async_train(rollout_id, rollout_data_curr_ref))
 
         if args.save_interval is not None and (
             (rollout_id + 1) % args.save_interval == 0