576 lines
23 KiB
Python
576 lines
23 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import logging
|
|
import os
|
|
import random
|
|
import socket
|
|
from argparse import Namespace
|
|
from contextlib import nullcontext
|
|
|
|
import ray
|
|
import torch
|
|
import torch.distributed as dist
|
|
from megatron.core import mpu
|
|
from ray.actor import ActorHandle
|
|
from torch_memory_saver import torch_memory_saver
|
|
from transformers import AutoConfig, AutoTokenizer
|
|
|
|
from slime.ray.train_actor import TrainRayActor
|
|
from slime.utils import train_dump_utils
|
|
from slime.utils.context_utils import with_defer
|
|
from slime.utils.data import process_rollout_data
|
|
from slime.utils.distributed_utils import get_gloo_group, init_process_group
|
|
from slime.utils.memory_utils import clear_memory, print_memory
|
|
from slime.utils.ray_utils import Box
|
|
from slime.utils.reloadable_process_group import destroy_process_groups, monkey_patch_torch_dist, reload_process_groups
|
|
from slime.utils.routing_replay import RoutingReplay
|
|
from slime.utils.timer import Timer, inverse_timer, timer
|
|
from slime.utils.tracking_utils import init_tracking
|
|
from slime.utils.types import RolloutBatch
|
|
|
|
from ...utils.profile_utils import TrainProfiler
|
|
from ...utils.tensor_backper import TensorBackuper
|
|
from .checkpoint import load_checkpoint
|
|
from .cp_utils import slice_log_prob_with_cp, slice_with_cp
|
|
from .data import DataIterator, get_data_iterator, log_perf_data, log_rollout_data, sync_actor_critic_data
|
|
from .initialize import init, is_megatron_main_rank
|
|
from .loss import compute_advantages_and_returns, get_log_probs_and_entropy, get_values
|
|
from .model import forward_only, initialize_model_and_optimizer, save, train
|
|
from .update_weight.common import named_params_and_buffers
|
|
from .update_weight.update_weight_from_distributed import UpdateWeightFromDistributed
|
|
from .update_weight.update_weight_from_tensor import UpdateWeightFromTensor
|
|
|
|
logging.getLogger("megatron").setLevel(logging.WARNING)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MegatronTrainRayActor(TrainRayActor):
|
|
@with_defer(lambda: Timer().start("train_wait"))
|
|
def init(
|
|
self,
|
|
args: Namespace,
|
|
role: str,
|
|
with_ref: bool = False,
|
|
) -> int | None:
|
|
monkey_patch_torch_dist()
|
|
|
|
super().init(args, role, with_ref)
|
|
|
|
init(args)
|
|
|
|
if is_megatron_main_rank():
|
|
init_tracking(args, primary=False)
|
|
|
|
self.prof = TrainProfiler(args)
|
|
|
|
# read config and tokenizer serialized to prevent concurrent writing bug.
|
|
for i in range(dist.get_world_size()):
|
|
if i == dist.get_rank():
|
|
self.hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True)
|
|
self.tokenizer = AutoTokenizer.from_pretrained(self.args.hf_checkpoint, trust_remote_code=True)
|
|
dist.barrier(group=get_gloo_group())
|
|
|
|
self.train_parallel_config = {
|
|
"dp_size": mpu.get_data_parallel_world_size(with_context_parallel=False),
|
|
}
|
|
dist.barrier(group=get_gloo_group())
|
|
|
|
if args.offload_train:
|
|
if (x := args.train_memory_margin_bytes) > 0:
|
|
logger.info(f"Set torch_memory_saver.memory_margin_bytes to {x}")
|
|
torch_memory_saver.memory_margin_bytes = x
|
|
|
|
if self.args.debug_rollout_only:
|
|
return 0
|
|
|
|
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.args.lr_warmup_iters = self.args.critic_lr_warmup_iters
|
|
|
|
(self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id) = initialize_model_and_optimizer(
|
|
args, role
|
|
)
|
|
|
|
if role == "critic":
|
|
if self.args.offload_train:
|
|
self.sleep()
|
|
return
|
|
|
|
start_rollout_id = loaded_rollout_id + 1
|
|
|
|
self.weights_backuper = TensorBackuper.create(
|
|
source_getter=lambda: named_params_and_buffers(
|
|
self.args,
|
|
self.model,
|
|
convert_to_global_name=args.megatron_to_hf_mode == "raw",
|
|
translate_gpu_to_cpu=not self.args.enable_weights_backuper,
|
|
),
|
|
single_tag=None if args.enable_weights_backuper else "actor",
|
|
)
|
|
self._active_model_tag: str | None = "actor"
|
|
self.weights_backuper.backup("actor")
|
|
|
|
if with_ref:
|
|
self.load_other_checkpoint("ref", args.ref_load)
|
|
|
|
if self.args.keep_old_actor:
|
|
# Load old_actor checkpoint
|
|
self.load_other_checkpoint("old_actor", args.load)
|
|
# Create rollout_actor as a copy of current actor
|
|
if args.update_weights_interval == 1:
|
|
self.weights_backuper.backup("rollout_actor")
|
|
|
|
if self.args.vocab_size is None:
|
|
self.args.vocab_size = self.tokenizer.vocab_size
|
|
|
|
update_weight_cls = UpdateWeightFromTensor if self.args.colocate else UpdateWeightFromDistributed
|
|
self.weight_updater = update_weight_cls(
|
|
self.args,
|
|
self.model,
|
|
weights_getter=lambda: self.weights_backuper.get("actor"),
|
|
model_name=type(self.hf_config).__name__.lower() if self.args.model_name is None else self.args.model_name,
|
|
quantization_config=getattr(self.hf_config, "quantization_config", None),
|
|
)
|
|
|
|
# empty cache after initialization
|
|
clear_memory()
|
|
|
|
if self.args.offload_train:
|
|
# recover to actor in the end.
|
|
self._switch_model("actor")
|
|
self.sleep()
|
|
|
|
self.rollout_engines = None
|
|
|
|
self.rollout_data_postprocess = None
|
|
if self.args.rollout_data_postprocess_path is not None:
|
|
from slime.utils.misc import load_function
|
|
|
|
self.rollout_data_postprocess = load_function(self.args.rollout_data_postprocess_path)
|
|
|
|
self.prof.on_init_end()
|
|
|
|
return start_rollout_id
|
|
|
|
@timer
|
|
def sleep(self) -> None:
|
|
assert self.args.offload_train
|
|
|
|
clear_memory(clear_host_memory=True)
|
|
print_memory("before offload model")
|
|
destroy_process_groups()
|
|
|
|
torch_memory_saver.pause()
|
|
|
|
print_memory("after offload model")
|
|
|
|
@timer
|
|
def wake_up(self) -> None:
|
|
assert self.args.offload_train
|
|
print_memory("before wake_up model")
|
|
|
|
torch_memory_saver.resume()
|
|
|
|
clear_memory()
|
|
reload_process_groups()
|
|
print_memory("after wake_up model")
|
|
|
|
def _get_rollout_data(self, rollout_data_ref: Box) -> RolloutBatch:
|
|
# Fetch data through ray on CPU, not sure if this will be performance bottleneck.
|
|
# Both first pp stage and the last pp stage will receive the data.
|
|
rollout_data = process_rollout_data(
|
|
self.args,
|
|
rollout_data_ref,
|
|
mpu.get_data_parallel_rank(with_context_parallel=False),
|
|
mpu.get_data_parallel_world_size(with_context_parallel=False),
|
|
)
|
|
# TODO: this is ugly, move to somewhere else?
|
|
# move tokens to GPU in advance
|
|
rollout_data["tokens"] = [
|
|
torch.tensor(t, dtype=torch.long, device=torch.cuda.current_device()) for t in rollout_data["tokens"]
|
|
]
|
|
rollout_data["loss_masks"] = [
|
|
torch.tensor(t, dtype=torch.int, device=torch.cuda.current_device()) for t in rollout_data["loss_masks"]
|
|
]
|
|
if "multimodal_train_inputs" in rollout_data:
|
|
# Move multimodal training tensors to GPU in advance
|
|
rollout_data["multimodal_train_inputs"] = [
|
|
(
|
|
{key: tensor.to(device=torch.cuda.current_device()) for key, tensor in mm_dict.items()}
|
|
if mm_dict is not None
|
|
else None
|
|
)
|
|
for mm_dict in rollout_data["multimodal_train_inputs"]
|
|
]
|
|
if "rollout_log_probs" in rollout_data:
|
|
rollout_data["rollout_log_probs"] = [
|
|
torch.tensor(
|
|
slice_log_prob_with_cp(log_prob, total_length, response_length),
|
|
device=torch.cuda.current_device(),
|
|
dtype=torch.float32,
|
|
)
|
|
for log_prob, total_length, response_length in zip(
|
|
rollout_data["rollout_log_probs"],
|
|
rollout_data["total_lengths"],
|
|
rollout_data["response_lengths"],
|
|
strict=False,
|
|
)
|
|
]
|
|
if "rollout_routed_experts" in rollout_data:
|
|
rollout_data["rollout_routed_experts"] = [
|
|
torch.from_numpy(r) for r in rollout_data["rollout_routed_experts"]
|
|
]
|
|
return rollout_data
|
|
|
|
def _switch_model(self, target_tag: str) -> None:
|
|
if target_tag not in self.weights_backuper.backup_tags:
|
|
raise ValueError(f"Cannot switch to unknown model tag: {target_tag}")
|
|
self.weights_backuper.restore(target_tag)
|
|
self._active_model_tag = target_tag
|
|
|
|
def fill_routing_replay(self, data_iterator, num_microbatches, rollout_data):
|
|
if "rollout_routed_experts" not in rollout_data:
|
|
raise ValueError(
|
|
"rollout_routed_experts is required in rollout_data when use_rollout_routing_replay is set."
|
|
)
|
|
|
|
from megatron.core.transformer.transformer_block import get_num_layers_to_build
|
|
from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
|
|
|
|
from slime.utils.routing_replay import RoutingReplay
|
|
|
|
for iterator in data_iterator:
|
|
iterator.reset()
|
|
|
|
tp_rank = mpu.get_tensor_model_parallel_rank()
|
|
tp_size = mpu.get_tensor_model_parallel_world_size()
|
|
|
|
def pad_func(experts, pad):
|
|
_, num_layers, topk = experts.shape
|
|
pad = (
|
|
torch.arange(
|
|
pad * num_layers * topk,
|
|
device=experts.device,
|
|
dtype=experts.dtype,
|
|
).reshape((pad, num_layers, topk))
|
|
% self.args.num_experts
|
|
)
|
|
return torch.cat([experts, pad], dim=0)
|
|
|
|
for _ in range(sum(num_microbatches)):
|
|
batch = data_iterator[0].get_next(["rollout_routed_experts", "tokens"])
|
|
rollout_routed_experts = batch["rollout_routed_experts"]
|
|
tokens = batch["tokens"]
|
|
assert len(rollout_routed_experts) == len(tokens)
|
|
for a, b in zip(rollout_routed_experts, tokens, strict=False):
|
|
assert a.shape[0] == b.shape[0] - 1, f"{a.shape}, {b.shape}"
|
|
|
|
# We need to pad the experts to the last token. We won't calculate loss on this token so this should be fine.
|
|
# TODO: fuse this padding with the following slice_with_cp to reduce memory copy.
|
|
rollout_routed_experts = [pad_func(r, 1) for r in rollout_routed_experts]
|
|
# TODO: maybe extract a common process function for here and get_batch?
|
|
rollout_routed_experts = [slice_with_cp(r, pad_func) for r in rollout_routed_experts]
|
|
rollout_routed_experts = torch.cat(rollout_routed_experts, dim=0)
|
|
pad_size = mpu.get_tensor_model_parallel_world_size() * self.args.data_pad_size_multiplier
|
|
pad = (pad_size - rollout_routed_experts.size(0) % pad_size) % pad_size
|
|
if pad != 0:
|
|
rollout_routed_experts = pad_func(rollout_routed_experts, pad)
|
|
|
|
if self.args.sequence_parallel:
|
|
seqlen = rollout_routed_experts.size(0)
|
|
assert seqlen % tp_size == 0
|
|
start, end = seqlen // tp_size * tp_rank, seqlen // tp_size * (tp_rank + 1)
|
|
rollout_routed_experts = rollout_routed_experts[start:end]
|
|
|
|
routing_replay_offset = 0
|
|
for vp_stage, model in enumerate(self.model):
|
|
config = model.module.config
|
|
num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage)
|
|
offset = get_transformer_layer_offset(config, vp_stage=vp_stage)
|
|
for layer_id in range(offset, offset + num_layers_to_build):
|
|
# skip dense layer
|
|
if isinstance(config.moe_layer_freq, int):
|
|
if layer_id % config.moe_layer_freq != 0:
|
|
continue
|
|
elif isinstance(config.moe_layer_freq, list):
|
|
assert len(config.moe_layer_freq) == config.num_layers
|
|
if config.moe_layer_freq[layer_id] == 0:
|
|
continue
|
|
layer_routed_experts = rollout_routed_experts[:, layer_id]
|
|
RoutingReplay.all_routing_replays[routing_replay_offset].record(layer_routed_experts)
|
|
routing_replay_offset += 1
|
|
assert routing_replay_offset == len(RoutingReplay.all_routing_replays)
|
|
|
|
del rollout_data["rollout_routed_experts"]
|
|
|
|
for iterator in data_iterator:
|
|
iterator.reset()
|
|
|
|
def compute_log_prob(
|
|
self,
|
|
data_iterator: list[DataIterator],
|
|
num_microbatches: list[int],
|
|
store_prefix: str = "",
|
|
) -> dict[str, list[torch.Tensor]]:
|
|
|
|
with timer(f"{store_prefix}log_probs"):
|
|
return forward_only(
|
|
get_log_probs_and_entropy,
|
|
self.args,
|
|
self.model,
|
|
data_iterator,
|
|
num_microbatches,
|
|
store_prefix=store_prefix,
|
|
)
|
|
|
|
def train(self, rollout_id: int, rollout_data_ref: Box) -> None:
|
|
if self.args.offload_train:
|
|
self.wake_up()
|
|
|
|
with timer("data_preprocess"):
|
|
rollout_data = self._get_rollout_data(rollout_data_ref)
|
|
if self.args.debug_rollout_only:
|
|
log_rollout_data(rollout_id, self.args, rollout_data)
|
|
return
|
|
|
|
if self.role == "critic":
|
|
return self.train_critic(rollout_id, rollout_data)
|
|
else:
|
|
return self.train_actor(rollout_id, rollout_data)
|
|
|
|
def train_critic(self, rollout_id: int, rollout_data: RolloutBatch) -> None:
|
|
# Create data iterator for log_probs and train.
|
|
data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data)
|
|
rollout_data.update(
|
|
forward_only(
|
|
get_values,
|
|
self.args,
|
|
self.model,
|
|
data_iterator,
|
|
num_microbatches,
|
|
)
|
|
)
|
|
|
|
if rollout_id >= self.args.num_critic_only_steps:
|
|
sync_actor_critic_data(self.args, rollout_data, self._actor_critic_groups)
|
|
|
|
compute_advantages_and_returns(self.args, rollout_data)
|
|
|
|
self.args.loss_type = "value_loss"
|
|
train(
|
|
rollout_id,
|
|
self.model,
|
|
self.optimizer,
|
|
self.opt_param_scheduler,
|
|
data_iterator,
|
|
num_microbatches,
|
|
)
|
|
|
|
def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None:
|
|
# Create data iterator for log_probs and train.
|
|
data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data)
|
|
|
|
if self.args.use_rollout_routing_replay:
|
|
self.fill_routing_replay(data_iterator, num_microbatches, rollout_data)
|
|
|
|
with inverse_timer("train_wait"), timer("train"):
|
|
if self.args.compute_advantages_and_returns:
|
|
if "ref" in self.weights_backuper.backup_tags:
|
|
if self.args.use_routing_replay:
|
|
os.environ["ROUTING_REPLAY_STAGE"] = "fallthrough"
|
|
self._switch_model("ref")
|
|
rollout_data.update(
|
|
self.compute_log_prob(
|
|
data_iterator,
|
|
num_microbatches,
|
|
store_prefix="ref_",
|
|
)
|
|
)
|
|
self._switch_model("old_actor" if self.args.keep_old_actor else "actor")
|
|
if not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics:
|
|
if self.args.use_routing_replay:
|
|
if self.args.use_rollout_routing_replay:
|
|
os.environ["ROUTING_REPLAY_STAGE"] = "replay_forward"
|
|
else:
|
|
os.environ["ROUTING_REPLAY_STAGE"] = "record"
|
|
rollout_data.update(
|
|
self.compute_log_prob(
|
|
data_iterator,
|
|
num_microbatches,
|
|
store_prefix="",
|
|
)
|
|
)
|
|
if self.args.use_rollout_routing_replay:
|
|
RoutingReplay.clear_all_forward()
|
|
|
|
if self.args.use_critic:
|
|
sync_actor_critic_data(
|
|
self.args,
|
|
rollout_data,
|
|
self._actor_critic_groups,
|
|
)
|
|
if self._active_model_tag != "actor":
|
|
self._switch_model("actor")
|
|
|
|
# Calculate adv and returns. Need to performed before training (instead of on the fly),
|
|
# because we may need normalize the whole rollout.
|
|
compute_advantages_and_returns(self.args, rollout_data)
|
|
|
|
if self.rollout_data_postprocess is not None:
|
|
self.rollout_data_postprocess(self.args)
|
|
|
|
log_rollout_data(rollout_id, self.args, rollout_data)
|
|
|
|
# Train
|
|
if self.args.use_routing_replay:
|
|
os.environ["ROUTING_REPLAY_STAGE"] = "replay_backward"
|
|
with timer("actor_train"):
|
|
train(
|
|
rollout_id,
|
|
self.model,
|
|
self.optimizer,
|
|
self.opt_param_scheduler,
|
|
data_iterator,
|
|
num_microbatches,
|
|
)
|
|
|
|
self.prof.step(rollout_id=rollout_id)
|
|
|
|
train_dump_utils.save_debug_train_data(self.args, rollout_id=rollout_id, rollout_data=rollout_data)
|
|
|
|
if self.args.use_routing_replay:
|
|
RoutingReplay.clear_all()
|
|
|
|
# update the cpu actor weight to the latest model
|
|
self.weights_backuper.backup("actor")
|
|
|
|
# Update ref model if needed
|
|
if (
|
|
self.args.ref_update_interval is not None
|
|
and (rollout_id + 1) % self.args.ref_update_interval == 0
|
|
and "ref" in self.weights_backuper.backup_tags
|
|
):
|
|
with timer("ref_model_update"):
|
|
if is_megatron_main_rank():
|
|
logger.info(f"Updating ref model at rollout_id {rollout_id}")
|
|
self.weights_backuper.backup("ref")
|
|
|
|
log_perf_data(rollout_id, self.args)
|
|
|
|
@timer
|
|
def save_model(self, rollout_id: int, force_sync: bool = False) -> None:
|
|
if self.args.debug_rollout_only:
|
|
return
|
|
|
|
# torch dist may trigger nccl communication during saving.
|
|
if self.args.offload_train:
|
|
reload_process_groups()
|
|
|
|
if self.args.async_save:
|
|
from megatron.training.async_utils import maybe_finalize_async_save
|
|
|
|
maybe_finalize_async_save(blocking=True)
|
|
|
|
save(rollout_id, self.model, self.optimizer, self.opt_param_scheduler)
|
|
|
|
if force_sync and self.args.async_save:
|
|
maybe_finalize_async_save(blocking=True)
|
|
|
|
if self.args.offload_train:
|
|
destroy_process_groups()
|
|
|
|
@timer
|
|
def update_weights(self) -> None:
|
|
if self.args.debug_train_only or self.args.debug_rollout_only:
|
|
return
|
|
|
|
if self.args.offload_train:
|
|
reload_process_groups()
|
|
|
|
rollout_engines, rollout_engine_lock, num_new_engines = ray.get(
|
|
self.rollout_manager.get_rollout_engines_and_lock.remote()
|
|
)
|
|
if num_new_engines > 0:
|
|
self.weight_updater.connect_rollout_engines(rollout_engines, rollout_engine_lock)
|
|
dist.barrier(group=get_gloo_group())
|
|
|
|
with torch_memory_saver.disable() if self.args.offload_train else nullcontext():
|
|
print_memory("before update_weights")
|
|
self.weight_updater.update_weights()
|
|
print_memory("after update_weights")
|
|
|
|
if self.args.ci_test and len(rollout_engines) > 0:
|
|
engine = random.choice(rollout_engines)
|
|
engine_version = ray.get(engine.get_weight_version.remote())
|
|
if str(engine_version) != str(self.weight_updater.weight_version):
|
|
raise RuntimeError(
|
|
f"Weight version mismatch! Engine: {engine_version}, Updater: {self.weight_updater.weight_version}"
|
|
)
|
|
|
|
if getattr(self.args, "keep_old_actor", False):
|
|
if self.args.update_weights_interval == 1:
|
|
logger.info("updating model queue: rollout_actor -> old_actor, actor -> rollout_actor")
|
|
# Queue-style update: rollout_actor params -> old_actor, actor params -> rollout_actor
|
|
# First copy rollout_actor to old_actor
|
|
self.weights_backuper.copy(src_tag="rollout_actor", dst_tag="old_actor")
|
|
# Then copy current actor to rollout_actor
|
|
self.weights_backuper.backup("rollout_actor")
|
|
else:
|
|
self.weights_backuper.backup("old_actor")
|
|
|
|
if self.args.offload_train:
|
|
destroy_process_groups()
|
|
|
|
def load_other_checkpoint(self, model_tag: str, path: str) -> None:
|
|
old_args = self.args.load, self.args.no_load_optim, self.args.no_load_rng, self.args.finetune
|
|
self.args.load = path
|
|
self.args.no_load_optim = True
|
|
self.args.no_load_rng = True
|
|
self.args.finetune = True
|
|
|
|
if model_tag == "ref" and self.args.ref_ckpt_step is not None:
|
|
old_ckpt_step = self.args.ckpt_step
|
|
self.args.ckpt_step = self.args.ref_ckpt_step
|
|
|
|
_, _ = load_checkpoint(
|
|
self.model,
|
|
None,
|
|
None,
|
|
checkpointing_context={},
|
|
skip_load_to_model_and_opt=False,
|
|
)
|
|
self.args.load, self.args.no_load_optim, self.args.no_load_rng, self.args.finetune = old_args
|
|
|
|
if model_tag == "ref" and self.args.ref_ckpt_step is not None:
|
|
self.args.ckpt_step = old_ckpt_step
|
|
|
|
self.weights_backuper.backup(model_tag)
|
|
self._active_model_tag = model_tag
|
|
|
|
def connect_actor_critic(
|
|
self,
|
|
actor_handle: ActorHandle | None = None,
|
|
master_address: str | None = None,
|
|
master_port: int | None = None,
|
|
) -> None:
|
|
if self.role == "actor":
|
|
master_address = ray.util.get_node_ip_address()
|
|
with socket.socket() as sock:
|
|
sock.bind(("", 0))
|
|
master_port = sock.getsockname()[1]
|
|
actor_handle.connect_actor_critic.remote(master_address=master_address, master_port=master_port)
|
|
|
|
group_name = "actor_critic"
|
|
world_size = 2
|
|
self._actor_critic_groups = init_process_group(
|
|
backend="nccl",
|
|
init_method=f"tcp://{master_address}:{master_port}",
|
|
world_size=world_size,
|
|
rank=0 if self.role == "actor" else 1,
|
|
group_name=group_name,
|
|
)
|