Revert "Rebase to latest slime"#11
Conversation
This reverts commit 2cb7c80.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new FSDP training backend and "True On-Policy" training capabilities, while performing a broad cleanup of legacy vLLM code and SGLang-specific naming. The updates include new FSDP utilities, Dockerfile refinements, and expanded documentation and examples. Reviewers identified several critical issues in the new FSDP modules, including a large number of missing imports, incorrect pathing in the build script, and multiple mathematical and logic errors in the PPO loss calculation, metric aggregation, and data unpacking utilities.
| import logging | ||
| import os | ||
| import random | ||
| from argparse import Namespace | ||
| from itertools import accumulate | ||
|
|
||
| import ray | ||
| import torch | ||
| import torch.distributed as dist | ||
| from tqdm import tqdm | ||
| from transformers import AutoConfig | ||
|
|
||
| from slime.ray.train_actor import TrainRayActor | ||
| from slime.utils import logging_utils, train_dump_utils, train_metric_utils | ||
| from slime.utils.data import get_minimum_num_micro_batch_size, process_rollout_data | ||
| from slime.utils.distributed_utils import get_gloo_group | ||
| from slime.utils.logging_utils import init_tracking | ||
| from slime.utils.memory_utils import clear_memory, print_memory | ||
| from slime.utils.metric_utils import compute_rollout_step | ||
| from slime.utils.misc import Box | ||
| from slime.utils.ppo_utils import compute_approx_kl, compute_gspo_kl, compute_opsm_mask, compute_policy_loss | ||
| from slime.utils.processing_utils import load_processor, load_tokenizer | ||
| from slime.utils.profile_utils import TrainProfiler | ||
| from slime.utils.timer import Timer, inverse_timer, timer, with_defer | ||
|
|
||
| from . import checkpoint | ||
| from .data_packing import pack_sequences, unpack_sequences | ||
| from .lr_scheduler import get_lr_scheduler | ||
| from .update_weight_utils import UpdateWeightFromDistributed, UpdateWeightFromTensor |
There was a problem hiding this comment.
The FSDPTrainRayActor implementation is missing several critical imports required for its methods to function. Specifically, the following symbols are used but not defined or imported:
sync_actor_critic_data(used at lines 384, 450)destroy_process_groups(used at lines 169, 532)reload_process_groups(used at lines 183, 514, 549)save_hf_model(used at line 529)ModelType(used at line 110)mpu(used at line 395)get_gpt_layer_with_transformer_engine_spec(used at line 52)IdentityOp(used at line 955)nullcontext(used at line 562)get_data_iterator(used at lines 372, 400)forward_only(used at line 375)get_values(used at line 376)compute_advantages_and_returns(used at lines 387, 461)log_rollout_data(used at line 466)get_grpo_returns(used at line 505)RolloutBatch(used as type hint at lines 352, 370, 398)
Please ensure all internal dependencies are correctly imported, likely from slime.backends.megatron_utils or slime.utils.
| import abc | ||
| import logging | ||
| import socket | ||
| from argparse import Namespace | ||
| from collections.abc import Sequence | ||
|
|
||
| import ray | ||
| import torch | ||
| import torch.distributed as dist | ||
| from ray.actor import ActorHandle | ||
| from torch.distributed.tensor import DTensor, Replicate | ||
|
|
||
| try: | ||
| from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions # type: ignore[import] | ||
| except ImportError: | ||
| from sglang.srt.patch_torch import monkey_patch_torch_reductions # type: ignore[import] | ||
|
|
||
| from sglang.srt.utils import MultiprocessingSerializer | ||
|
|
||
| from slime.utils.distributed_utils import init_process_group | ||
|
|
||
|
|
||
| try: | ||
| from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket # type: ignore[import] | ||
| except ImportError: | ||
| from sglang.srt.model_executor.model_runner import FlattenedTensorBucket # type: ignore[import] | ||
|
|
There was a problem hiding this comment.
Missing critical imports in update_weight_utils.py:
update_weights_from_distributed(used at line 267) is not imported.ObjectRef(used at line 260) is not imported fromray.tqdm(used at lines 104, 124) is not imported.
These missing symbols will cause runtime crashes during weight synchronization.
| else | ||
| export SLIME_DIR=$BASE_DIR/slime | ||
| cd $SLIME_DIR | ||
| export SLIME_DIR=$BASE_DIR/ |
There was a problem hiding this comment.
The SLIME_DIR is set incorrectly in the else block. It should point to the repository directory (e.g., $BASE_DIR/slime) rather than the parent directory. This will cause subsequent commands that rely on $SLIME_DIR (like the patch applications on lines 81 and 83) to fail because they won't find the docker/ directory.
| export SLIME_DIR=$BASE_DIR/ | |
| export SLIME_DIR=$BASE_DIR/slime |
|
|
||
| advantages = advantages.to(device=log_probs.device) | ||
| old_log_probs = old_log_probs.to(device=log_probs.device) | ||
| ppo_kl = old_log_probs - log_probs |
There was a problem hiding this comment.
The calculation of ppo_kl as old_log_probs - log_probs results in the negative log-ratio. If compute_policy_loss expects the standard log-ratio (new - old) to compute the PPO ratio exp(new - old), this will result in an inverted gradient, causing the model to minimize rather than maximize the expected reward.
| ppo_kl = old_log_probs - log_probs | |
| ppo_kl = log_probs - old_log_probs |
| dist.all_gather_object(reduced_aggregated, aggregated, group=self.dp_group) | ||
| aggregated = {} | ||
| for k in reported_accum.keys(): | ||
| aggregated[k] = sum([r[k] for r in reduced_aggregated]) / (self.args.global_batch_size) |
There was a problem hiding this comment.
The metric aggregation logic is incorrect. aggregated[k] already contains the sum of scaled losses (scaled by dp_size / global_batch_size) for the local rank. Summing these across all ranks gives (dp_size / global_batch_size) * total_global_loss. Dividing by global_batch_size again results in an incorrect value. It should be divided by self.dp_size to recover the global average.
| aggregated[k] = sum([r[k] for r in reduced_aggregated]) / (self.args.global_batch_size) | |
| aggregated[k] = sum([r[k] for r in reduced_aggregated]) / self.dp_size |
| ] | ||
| elif key == "rollout_log_probs": | ||
| # rollout_log_probs is packed based on response_lengths, so slice differently |
There was a problem hiding this comment.
The slicing logic in unpack_sequences incorrectly applies pad_length to every sequence in the batch. Since padding is typically added only at the end of the concatenated batch, subtracting pad_length from the indices of intermediate sequences will result in incorrect slices. The indices in cu_seqlens are already relative to the start of the valid data, so pad_length should not be used in the slice calculation for individual sequences.
if key in ["log_probs", "ref_log_probs", "cur_log_probs", "entropy"]:
# These are computed from logits[:-1] so they have length seq_len-1
instance[key] = value[
end_idx - 1 - response_lengths[i] : end_idx - 1
]| log_probs_full = torch.log_softmax(shifted_logits, dim=-1) | ||
| probs = torch.softmax(shifted_logits, dim=-1) | ||
| entropy = -(probs * log_probs_full).sum(dim=-1) | ||
| return log_probs, entropy |
There was a problem hiding this comment.
The entropy calculation does not account for the temperature parameter. While gather_log_probs_packed correctly applies the temperature scaling to the logits, the subsequent log_softmax and softmax calls for entropy use the unscaled shifted_logits. This leads to a mathematical inconsistency between the policy loss and the entropy bonus.
1 participant
Reverts #10