731 lines
29 KiB
Python
731 lines
29 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import dataclasses
|
|
import gc
|
|
import logging
|
|
import math
|
|
import os
|
|
from argparse import Namespace
|
|
from collections.abc import Callable, Sequence
|
|
from functools import partial
|
|
|
|
import torch
|
|
from megatron.core import mpu
|
|
from megatron.core.distributed import DistributedDataParallel as DDP
|
|
from megatron.core.distributed import finalize_model_grads
|
|
from megatron.core.enums import ModelType
|
|
from megatron.core.models.gpt import GPTModel
|
|
from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer
|
|
from megatron.core.optimizer.optimizer import MegatronOptimizer
|
|
from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
|
|
from megatron.core.pipeline_parallel import get_forward_backward_func
|
|
from megatron.core.utils import get_model_config
|
|
from megatron.training.global_vars import get_args
|
|
from megatron.training.training import get_model
|
|
|
|
from slime.utils import tracking_utils
|
|
from slime.utils.memory_utils import clear_memory
|
|
|
|
from .checkpoint import load_checkpoint, save_checkpoint
|
|
from .data import DataIterator, get_batch
|
|
from .loss import loss_function
|
|
from .model_provider import get_model_provider_func
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def get_optimizer_param_scheduler(args: Namespace, optimizer: MegatronOptimizer) -> OptimizerParamScheduler:
|
|
"""Create and configure the optimizer learning-rate/weight-decay scheduler.
|
|
|
|
This configures iteration-based schedules derived from the global batch size
|
|
and run-time arguments.
|
|
|
|
Args:
|
|
args (Namespace): Training/runtime arguments (argparse namespace).
|
|
optimizer (MegatronOptimizer): Megatron optimizer bound to the model.
|
|
|
|
Returns:
|
|
OptimizerParamScheduler: Initialized scheduler bound to ``optimizer``.
|
|
"""
|
|
# Iteration-based training.
|
|
args.train_iters = args.num_rollout * args.rollout_batch_size * args.n_samples_per_prompt // args.global_batch_size
|
|
if args.lr_decay_iters is None:
|
|
args.lr_decay_iters = args.train_iters
|
|
lr_decay_steps = args.lr_decay_iters * args.global_batch_size
|
|
wd_incr_steps = args.train_iters * args.global_batch_size
|
|
wsd_decay_steps = None
|
|
if args.lr_wsd_decay_iters is not None:
|
|
wsd_decay_steps = args.lr_wsd_decay_iters * args.global_batch_size
|
|
if args.lr_warmup_fraction is not None:
|
|
lr_warmup_steps = args.lr_warmup_fraction * lr_decay_steps
|
|
else:
|
|
lr_warmup_steps = args.lr_warmup_iters * args.global_batch_size
|
|
|
|
opt_param_scheduler = OptimizerParamScheduler(
|
|
optimizer,
|
|
init_lr=args.lr_warmup_init,
|
|
max_lr=args.lr,
|
|
min_lr=args.min_lr,
|
|
lr_warmup_steps=lr_warmup_steps,
|
|
lr_decay_steps=lr_decay_steps,
|
|
lr_decay_style=args.lr_decay_style,
|
|
start_wd=args.start_weight_decay,
|
|
end_wd=args.end_weight_decay,
|
|
wd_incr_steps=wd_incr_steps,
|
|
wd_incr_style=args.weight_decay_incr_style,
|
|
use_checkpoint_opt_param_scheduler=args.use_checkpoint_opt_param_scheduler,
|
|
override_opt_param_scheduler=args.override_opt_param_scheduler,
|
|
wsd_decay_steps=wsd_decay_steps,
|
|
lr_wsd_decay_style=args.lr_wsd_decay_style,
|
|
)
|
|
|
|
return opt_param_scheduler
|
|
|
|
|
|
def setup_model_and_optimizer(
|
|
args: Namespace,
|
|
role: str = "actor",
|
|
) -> tuple[list[DDP], MegatronOptimizer, OptimizerParamScheduler]:
|
|
"""Build model(s), wrap with DDP, and construct optimizer and scheduler.
|
|
|
|
Args:
|
|
args (Namespace): Training/runtime arguments (argparse namespace).
|
|
role (str): Logical role of the model (e.g., "actor", "critic").
|
|
no_wd_decay_cond (Callable[..., bool] | None): Predicate to exclude
|
|
parameters from weight decay.
|
|
scale_lr_cond (Callable[..., bool] | None): Predicate to scale LR for
|
|
selected parameter groups.
|
|
lr_mult (float): Global learning-rate multiplier for the optimizer.
|
|
|
|
Returns:
|
|
tuple[list[DDP], MegatronOptimizer, OptimizerParamScheduler]:
|
|
- List of model chunks wrapped by ``DDP``.
|
|
- The constructed ``MegatronOptimizer`` instance.
|
|
- The learning-rate/weight-decay scheduler tied to the optimizer.
|
|
"""
|
|
assert not args.moe_use_upcycling
|
|
assert args.load is not None or args.pretrained_checkpoint is not None
|
|
|
|
model = get_model(get_model_provider_func(args, role), ModelType.encoder_or_decoder)
|
|
|
|
# Optimizer
|
|
kwargs = {}
|
|
for f in dataclasses.fields(OptimizerConfig):
|
|
if hasattr(args, f.name):
|
|
kwargs[f.name] = getattr(args, f.name)
|
|
config = OptimizerConfig(**kwargs)
|
|
config.timers = None
|
|
|
|
optimizer = get_megatron_optimizer(
|
|
config=config,
|
|
model_chunks=model,
|
|
use_gloo_process_groups=args.enable_gloo_process_groups,
|
|
)
|
|
opt_param_scheduler = get_optimizer_param_scheduler(args, optimizer)
|
|
return model, optimizer, opt_param_scheduler
|
|
|
|
|
|
def enable_forward_pre_hook(model_chunks: Sequence[DDP]) -> None:
|
|
"""Enable forward pre-hooks for provided DDP-wrapped model chunks.
|
|
|
|
Args:
|
|
model_chunks (Sequence[DDP]): Sequence of DDP modules to enable hooks on.
|
|
"""
|
|
for model_chunk in model_chunks:
|
|
assert isinstance(model_chunk, DDP)
|
|
model_chunk.enable_forward_pre_hook()
|
|
|
|
|
|
def disable_forward_pre_hook(model_chunks: Sequence[DDP], param_sync: bool = True) -> None:
|
|
"""Disable forward pre-hooks for provided DDP-wrapped model chunks.
|
|
|
|
Args:
|
|
model_chunks (Sequence[DDP]): Sequence of DDP modules to disable hooks on.
|
|
param_sync (bool): Whether to synchronize parameters when disabling.
|
|
"""
|
|
for model_chunk in model_chunks:
|
|
assert isinstance(model_chunk, DDP)
|
|
model_chunk.disable_forward_pre_hook(param_sync=param_sync)
|
|
|
|
|
|
@torch.no_grad()
|
|
def forward_only(
|
|
f: Callable[..., dict[str, list[torch.Tensor]]],
|
|
args: Namespace,
|
|
model: Sequence[DDP],
|
|
data_iterator: Sequence[DataIterator],
|
|
num_microbatches: Sequence[int],
|
|
store_prefix: str = "",
|
|
) -> dict[str, list[torch.Tensor]]:
|
|
"""Run forward passes only and collect non-loss outputs (e.g., logprobs).
|
|
|
|
The model is put into evaluation mode, a forward-only pipeline pass is
|
|
executed, and relevant outputs are aggregated and returned.
|
|
|
|
Args:
|
|
f (Callable[..., dict[str, list[torch.Tensor]]]): Post-forward callback used to
|
|
compute and package outputs to collect. This should accept a logits
|
|
tensor as its first positional argument and additional keyword-only
|
|
arguments; see ``get_log_probs_and_entropy``/``get_values`` in
|
|
``megatron_utils.loss`` for examples. It will be partially applied
|
|
so that the callable returned from the internal forward step only
|
|
requires the logits tensor.
|
|
args (Namespace): Runtime arguments.
|
|
model (Sequence[DDP]): Sequence of DDP-wrapped model chunks.
|
|
data_iterator (Sequence[DataIterator]): Iterable(s) yielding batches for inference.
|
|
num_microbatches (Sequence[int]): Number of microbatches per rollout step.
|
|
store_prefix (str): Prefix to prepend to stored output keys.
|
|
|
|
Returns:
|
|
dict[str, list[torch.Tensor]]: Aggregated outputs keyed by ``store_prefix + key``.
|
|
"""
|
|
|
|
# reset data iterator
|
|
for iterator in data_iterator:
|
|
iterator.reset()
|
|
|
|
config = get_model_config(model[0])
|
|
|
|
def forward_step(
|
|
data_iterator: DataIterator, model: GPTModel, return_schedule_plan: bool = False
|
|
) -> tuple[torch.Tensor, Callable[[torch.Tensor], dict[str, list[torch.Tensor]]]]:
|
|
"""Forward step used by Megatron's pipeline engine.
|
|
|
|
Args:
|
|
data_iterator (DataIterator): Input data iterator.
|
|
model (GPTModel): The GPT model chunk to execute.
|
|
|
|
Returns:
|
|
tuple[torch.Tensor, Callable[[torch.Tensor], dict[str, list[torch.Tensor]]]]:
|
|
Output tensor(s) and a callable that computes and packages results
|
|
to be collected by the engine.
|
|
"""
|
|
|
|
assert not return_schedule_plan, "forward_only step should never return schedule plan"
|
|
|
|
# Get the batch.
|
|
batch = get_batch(
|
|
data_iterator,
|
|
[
|
|
"tokens",
|
|
"loss_masks",
|
|
"multimodal_train_inputs",
|
|
"total_lengths",
|
|
"response_lengths",
|
|
],
|
|
args.data_pad_size_multiplier,
|
|
)
|
|
unconcat_tokens = batch["unconcat_tokens"]
|
|
tokens = batch["tokens"]
|
|
packed_seq_params = batch["packed_seq_params"]
|
|
total_lengths = batch["total_lengths"]
|
|
response_lengths = batch["response_lengths"]
|
|
output_tensor = model(
|
|
input_ids=tokens,
|
|
position_ids=None,
|
|
attention_mask=None,
|
|
labels=None,
|
|
packed_seq_params=packed_seq_params,
|
|
loss_mask=batch["full_loss_masks"],
|
|
**(batch["multimodal_train_inputs"] if batch["multimodal_train_inputs"] is not None else {}),
|
|
)
|
|
|
|
return output_tensor, partial(
|
|
f,
|
|
args=args,
|
|
unconcat_tokens=unconcat_tokens,
|
|
total_lengths=total_lengths,
|
|
response_lengths=response_lengths,
|
|
with_entropy=args.use_rollout_entropy,
|
|
)
|
|
|
|
# Turn on evaluation mode which disables dropout.
|
|
for model_module in model:
|
|
model_module.eval()
|
|
|
|
if args.custom_megatron_before_log_prob_hook_path:
|
|
from slime.utils.misc import load_function
|
|
|
|
custom_before_log_prob_hook = load_function(args.custom_megatron_before_log_prob_hook_path)
|
|
custom_before_log_prob_hook(args, model, store_prefix)
|
|
|
|
forward_backward_func = get_forward_backward_func()
|
|
# Don't care about timing during evaluation
|
|
config.timers = None
|
|
forward_data_store = []
|
|
num_steps_per_rollout = len(num_microbatches)
|
|
for step_id in range(num_steps_per_rollout):
|
|
# collect_non_loss_data
|
|
forward_data_store += forward_backward_func(
|
|
forward_step_func=forward_step,
|
|
data_iterator=data_iterator,
|
|
model=model,
|
|
num_microbatches=num_microbatches[step_id],
|
|
seq_length=args.seq_length,
|
|
micro_batch_size=args.micro_batch_size,
|
|
forward_only=True,
|
|
collect_non_loss_data=True,
|
|
)
|
|
|
|
# Move model back to the train mode.
|
|
for model_module in model:
|
|
model_module.train()
|
|
|
|
rollout_data = {}
|
|
# Store the results on the last stage
|
|
if mpu.is_pipeline_last_stage():
|
|
keys = forward_data_store[0].keys()
|
|
for key in keys:
|
|
values = []
|
|
for value in forward_data_store:
|
|
assert isinstance(value[key], list)
|
|
values += value[key]
|
|
|
|
if args.use_dynamic_batch_size:
|
|
# TODO: This is ugly... Find a better way to make the data have the same order.
|
|
# TODO: move this out of the loop.
|
|
origin_values = [None] * len(values)
|
|
origin_indices = sum(data_iterator[0].micro_batch_indices, [])
|
|
for value, origin_index in zip(values, origin_indices, strict=False):
|
|
origin_values[origin_index] = value
|
|
values = origin_values
|
|
rollout_data[f"{store_prefix}{key}"] = values
|
|
return rollout_data
|
|
|
|
|
|
def train_one_step(
|
|
args: Namespace,
|
|
rollout_id: int,
|
|
step_id: int,
|
|
data_iterator: Sequence[DataIterator],
|
|
model: Sequence[DDP],
|
|
optimizer: MegatronOptimizer,
|
|
opt_param_scheduler: OptimizerParamScheduler,
|
|
num_microbatches: int,
|
|
) -> tuple[dict[str, float], float]:
|
|
"""Execute a single pipeline-parallel training step.
|
|
|
|
Runs forward/backward over ``num_microbatches``, applies optimizer step and
|
|
one scheduler step when gradients are valid.
|
|
|
|
Args:
|
|
args (Namespace): Runtime arguments.
|
|
rollout_id (int): Rollout identifier.
|
|
step_id (int): Step index within the current rollout.
|
|
data_iterator (Sequence[DataIterator]): Iterable(s) yielding training batches.
|
|
model (Sequence[DDP]): Sequence of DDP-wrapped model chunks.
|
|
optimizer (MegatronOptimizer): Optimizer instance.
|
|
opt_param_scheduler (OptimizerParamScheduler): LR/WD scheduler.
|
|
num_microbatches (int): Number of microbatches to process.
|
|
|
|
Returns:
|
|
tuple[dict[str, float], float]: Reduced loss dictionary (last stage only)
|
|
and gradient norm for logging.
|
|
"""
|
|
args = get_args()
|
|
|
|
# Set grad to zero.
|
|
for model_chunk in model:
|
|
model_chunk.zero_grad_buffer()
|
|
optimizer.zero_grad()
|
|
|
|
if args.custom_megatron_before_train_step_hook_path:
|
|
from slime.utils.misc import load_function
|
|
|
|
custom_before_train_step_hook = load_function(args.custom_megatron_before_train_step_hook_path)
|
|
custom_before_train_step_hook(args, rollout_id, step_id, model, optimizer, opt_param_scheduler)
|
|
|
|
def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_plan: bool = False) -> tuple[
|
|
torch.Tensor,
|
|
Callable[[torch.Tensor], tuple[torch.Tensor, int, dict[str, torch.Tensor | list[str]]]],
|
|
]:
|
|
"""Forward step used by Megatron's pipeline engine during training.
|
|
|
|
Args:
|
|
data_iterator (DataIterator): Input data iterator.
|
|
model (GPTModel): The GPT model chunk to execute.
|
|
|
|
Returns:
|
|
tuple[torch.Tensor, Callable[[torch.Tensor], tuple[torch.Tensor, int, dict[str, torch.Tensor | list[str]]]]]:
|
|
Output tensor(s) and the loss function, which returns
|
|
(loss, num_elems, {"keys": list[str], "values": torch.Tensor}).
|
|
"""
|
|
|
|
# Get the batch.
|
|
batch = get_batch(
|
|
data_iterator,
|
|
[
|
|
"tokens",
|
|
"multimodal_train_inputs",
|
|
"packed_seq_params",
|
|
"total_lengths",
|
|
"response_lengths",
|
|
"loss_masks",
|
|
"log_probs",
|
|
"ref_log_probs",
|
|
"values",
|
|
"advantages",
|
|
"returns",
|
|
"rollout_log_probs",
|
|
"teacher_log_probs", # For OPD distillation loss
|
|
],
|
|
args.data_pad_size_multiplier,
|
|
)
|
|
|
|
if os.environ.get("ENABLE_ROUTING_REPLAY", "0") == "1":
|
|
old_stage = os.environ["ROUTING_REPLAY_STAGE"]
|
|
os.environ["ROUTING_REPLAY_STAGE"] = "replay_forward"
|
|
|
|
if return_schedule_plan:
|
|
assert not args.enable_mtp_training, "MTP training should not be enabled when using combined 1f1b"
|
|
output_tensor = model.build_schedule_plan(
|
|
input_ids=batch["tokens"],
|
|
position_ids=None,
|
|
attention_mask=None,
|
|
labels=None,
|
|
packed_seq_params=batch["packed_seq_params"],
|
|
loss_mask=batch["full_loss_masks"],
|
|
)
|
|
else:
|
|
output_tensor = model(
|
|
input_ids=batch["tokens"],
|
|
position_ids=None,
|
|
attention_mask=None,
|
|
labels=None,
|
|
packed_seq_params=batch["packed_seq_params"],
|
|
loss_mask=batch["full_loss_masks"],
|
|
mtp_kwargs={"mtp_labels": batch["tokens"]} if args.enable_mtp_training else {},
|
|
**(batch["multimodal_train_inputs"] if batch["multimodal_train_inputs"] is not None else {}),
|
|
)
|
|
|
|
if os.environ.get("ENABLE_ROUTING_REPLAY", "0") == "1":
|
|
os.environ["ROUTING_REPLAY_STAGE"] = old_stage
|
|
|
|
return output_tensor, partial(loss_function, args, batch, num_microbatches)
|
|
|
|
# Forward pass.
|
|
forward_backward_func = get_forward_backward_func()
|
|
losses_reduced = forward_backward_func(
|
|
forward_step_func=forward_step,
|
|
data_iterator=data_iterator,
|
|
model=model,
|
|
num_microbatches=num_microbatches,
|
|
seq_length=args.seq_length,
|
|
micro_batch_size=args.micro_batch_size,
|
|
decoder_seq_length=args.decoder_seq_length,
|
|
forward_only=False,
|
|
)
|
|
|
|
valid_step = True
|
|
if not getattr(args, "check_for_nan_in_loss_and_grad", True):
|
|
found_inf_flag = optimizer.prepare_grads()
|
|
if found_inf_flag:
|
|
valid_step = False
|
|
else:
|
|
grad_norm = optimizer.get_grad_norm()
|
|
if isinstance(grad_norm, torch.Tensor):
|
|
valid_step = not (torch.isnan(grad_norm) or torch.isinf(grad_norm))
|
|
else:
|
|
valid_step = not (math.isnan(grad_norm) or math.isinf(grad_norm))
|
|
|
|
# CI check: verify only MTP parameters have non-zero gradients when truncation happens
|
|
# This check must happen before optimizer.step() as gradients may be modified during step
|
|
if args.ci_test and args.enable_mtp_training:
|
|
from slime.backends.megatron_utils.ci_utils import check_mtp_only_grad
|
|
|
|
check_mtp_only_grad(model, step_id)
|
|
|
|
if valid_step:
|
|
# Update parameters.
|
|
update_successful, grad_norm, num_zeros_in_grad = optimizer.step()
|
|
|
|
# Update learning rate.
|
|
assert update_successful
|
|
opt_param_scheduler.step(increment=args.global_batch_size)
|
|
|
|
# release grad
|
|
for model_chunk in model:
|
|
model_chunk.zero_grad_buffer()
|
|
optimizer.zero_grad()
|
|
|
|
if mpu.is_pipeline_last_stage(ignore_virtual=True):
|
|
# Average loss across microbatches.
|
|
keys = losses_reduced[0]["keys"]
|
|
values = None
|
|
for x in losses_reduced:
|
|
if values is None:
|
|
values = x["values"]
|
|
else:
|
|
values += x["values"]
|
|
assert len(keys) + 1 == values.numel()
|
|
torch.distributed.all_reduce(values, group=mpu.get_data_parallel_group(with_context_parallel=True))
|
|
|
|
loss_reduced = {}
|
|
values = values.tolist()
|
|
num_samples_or_tokens = values[0]
|
|
for key, value in zip(keys, values[1:], strict=False):
|
|
loss_reduced[key] = value * mpu.get_context_parallel_world_size() / num_samples_or_tokens
|
|
return loss_reduced, grad_norm
|
|
return {}, grad_norm
|
|
|
|
|
|
def should_disable_forward_pre_hook(args: Namespace) -> bool:
|
|
"""Block forward pre-hook for certain configurations."""
|
|
return args.use_distributed_optimizer and args.overlap_param_gather
|
|
|
|
|
|
def finalize_model_grads_with_empty_cache(*args, **kwargs):
|
|
# trigger empty cache when there are less than 10% free memory before the final reduce scatter.
|
|
# TODO: this is an ad-hoc method and we should figure out why the oom happens in the first place.
|
|
device = torch.cuda.current_device()
|
|
free, total = torch.cuda.mem_get_info(device)
|
|
if free / total < 0.1:
|
|
clear_memory()
|
|
return finalize_model_grads(*args, **kwargs)
|
|
|
|
|
|
def train(
|
|
rollout_id: int,
|
|
model: Sequence[DDP],
|
|
optimizer: MegatronOptimizer,
|
|
opt_param_scheduler: OptimizerParamScheduler,
|
|
data_iterator: Sequence[DataIterator],
|
|
num_microbatches: Sequence[int],
|
|
) -> None:
|
|
"""Run training over a rollout consisting of multiple steps.
|
|
|
|
The model is switched to train mode, training hooks are configured, and
|
|
``train_one_step`` is invoked for each step in the rollout.
|
|
|
|
Args:
|
|
rollout_id (int): Rollout identifier.
|
|
model (Sequence[DDP]): Sequence of DDP-wrapped model chunks.
|
|
optimizer (MegatronOptimizer): Optimizer instance.
|
|
opt_param_scheduler (OptimizerParamScheduler): LR/WD scheduler.
|
|
data_iterator (Sequence[DataIterator]): Iterable(s) yielding training batches.
|
|
num_microbatches (Sequence[int]): Microbatches per step in the rollout.
|
|
"""
|
|
args = get_args()
|
|
|
|
for iterator in data_iterator:
|
|
iterator.reset()
|
|
|
|
# Turn on training mode which enables dropout.
|
|
for model_module in model:
|
|
model_module.train()
|
|
|
|
# Setup some training config params.
|
|
config = get_model_config(model[0])
|
|
config.grad_scale_func = optimizer.scale_loss
|
|
config.timers = None
|
|
if isinstance(model[0], DDP) and args.overlap_grad_reduce:
|
|
assert config.no_sync_func is None, (
|
|
"When overlap_grad_reduce is True, config.no_sync_func must be None; "
|
|
"a custom no_sync_func is not supported when overlapping grad-reduce"
|
|
)
|
|
config.no_sync_func = [model_chunk.no_sync for model_chunk in model]
|
|
if len(model) == 1:
|
|
config.no_sync_func = config.no_sync_func[0]
|
|
if args.align_grad_reduce:
|
|
config.grad_sync_func = [model_chunk.start_grad_sync for model_chunk in model]
|
|
if len(model) == 1:
|
|
config.grad_sync_func = config.grad_sync_func[0]
|
|
if args.overlap_param_gather and args.align_param_gather:
|
|
config.param_sync_func = [model_chunk.start_param_sync for model_chunk in model]
|
|
if len(model) == 1:
|
|
config.param_sync_func = config.param_sync_func[0]
|
|
config.finalize_model_grads_func = finalize_model_grads_with_empty_cache
|
|
|
|
pre_hook_enabled = False
|
|
|
|
if args.manual_gc:
|
|
# Disable the default garbage collector and perform the collection manually.
|
|
# This is to align the timing of garbage collection across ranks.
|
|
assert args.manual_gc_interval >= 0, "Manual garbage collection interval should be larger than or equal to 0"
|
|
gc.disable()
|
|
gc.collect()
|
|
|
|
# Disable forward pre-hook to start training to ensure that errors in checkpoint loading
|
|
# or random initialization don't propagate to all ranks in first all-gather (which is a
|
|
# no-op if things work correctly).
|
|
if should_disable_forward_pre_hook(args):
|
|
disable_forward_pre_hook(model, param_sync=False)
|
|
# Also remove param_sync_func temporarily so that sync calls made in
|
|
# `forward_backward_func` are no-ops.
|
|
param_sync_func = config.param_sync_func
|
|
config.param_sync_func = None
|
|
pre_hook_enabled = False
|
|
|
|
num_steps_per_rollout = len(num_microbatches)
|
|
|
|
# Run training iterations till done.
|
|
for step_id in range(num_steps_per_rollout):
|
|
|
|
# Run training step.
|
|
loss_dict, grad_norm = train_one_step(
|
|
args,
|
|
rollout_id,
|
|
step_id,
|
|
data_iterator,
|
|
model,
|
|
optimizer,
|
|
opt_param_scheduler,
|
|
num_microbatches[step_id],
|
|
)
|
|
|
|
if step_id == 0:
|
|
# Enable forward pre-hook after training step has successfully run. All subsequent
|
|
# forward passes will use the forward pre-hook / `param_sync_func` in
|
|
# `forward_backward_func`.
|
|
if should_disable_forward_pre_hook(args):
|
|
enable_forward_pre_hook(model)
|
|
config.param_sync_func = param_sync_func
|
|
pre_hook_enabled = True
|
|
|
|
if args.enable_mtp_training:
|
|
from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper
|
|
|
|
mtp_loss_scale = 1 / num_microbatches[step_id]
|
|
tracker = MTPLossLoggingHelper.tracker
|
|
if "values" in tracker:
|
|
values = tracker["values"]
|
|
if tracker.get("reduce_group") is not None:
|
|
torch.distributed.all_reduce(values, group=tracker.get("reduce_group"))
|
|
if tracker.get("avg_group") is not None:
|
|
torch.distributed.all_reduce(values, group=tracker["avg_group"], op=torch.distributed.ReduceOp.AVG)
|
|
# here we assume only one mtp layer
|
|
mtp_losses = (tracker["values"] * mtp_loss_scale).item()
|
|
MTPLossLoggingHelper.clean_loss_in_tracker()
|
|
|
|
# CI check: verify MTP loss is within expected bounds
|
|
if args.ci_test:
|
|
from slime.backends.megatron_utils.ci_utils import check_mtp_loss
|
|
|
|
check_mtp_loss(mtp_losses)
|
|
|
|
# per train step log.
|
|
if (
|
|
mpu.get_data_parallel_rank(with_context_parallel=True) == 0
|
|
and mpu.get_tensor_model_parallel_rank() == 0
|
|
and mpu.get_pipeline_model_parallel_rank() == mpu.get_pipeline_model_parallel_world_size() - 1
|
|
):
|
|
accumulated_step_id = rollout_id * num_steps_per_rollout + step_id
|
|
role = getattr(model[0], "role", "actor")
|
|
role_tag = "" if role == "actor" else f"{role}-"
|
|
log_dict = {
|
|
f"train/{role_tag}{key}": val.mean().item() if isinstance(val, torch.Tensor) else val
|
|
for key, val in loss_dict.items()
|
|
}
|
|
log_dict[f"train/{role_tag}grad_norm"] = grad_norm
|
|
if args.enable_mtp_training:
|
|
log_dict[f"train/{role_tag}mtp_loss"] = mtp_losses
|
|
|
|
for param_group_id, param_group in enumerate(optimizer.param_groups):
|
|
log_dict[f"train/{role_tag}lr-pg_{param_group_id}"] = opt_param_scheduler.get_lr(param_group)
|
|
|
|
log_dict["train/step"] = accumulated_step_id
|
|
tracking_utils.log(args, log_dict, step_key="train/step")
|
|
|
|
if args.ci_test and not args.ci_disable_kl_checker:
|
|
if step_id == 0 and "train/ppo_kl" in log_dict and "train/pg_clipfrac" in log_dict:
|
|
if args.multi_latent_attention:
|
|
# TODO: mla currently have non-zero kl, need further investigation
|
|
assert log_dict["train/ppo_kl"] < 1e-8, f"{log_dict=}"
|
|
else:
|
|
assert log_dict["train/ppo_kl"] == 0.0 and log_dict["train/pg_clipfrac"] == 0.0, f"{log_dict=}"
|
|
if accumulated_step_id == 0 and "train/kl_loss" in log_dict:
|
|
assert log_dict["train/kl_loss"] == 0.0, f"{log_dict=}"
|
|
|
|
logger.info(f"{role_tag}step {accumulated_step_id}: {log_dict}")
|
|
|
|
if args.ci_save_grad_norm is not None:
|
|
ci_save_grad_norm_path = args.ci_save_grad_norm.format(
|
|
role=role,
|
|
rollout_id=rollout_id,
|
|
step_id=step_id,
|
|
)
|
|
torch.save(grad_norm, ci_save_grad_norm_path)
|
|
elif args.ci_load_grad_norm is not None:
|
|
ci_load_grad_norm_path = args.ci_load_grad_norm.format(
|
|
role=role,
|
|
rollout_id=rollout_id,
|
|
step_id=step_id,
|
|
)
|
|
expected_grad_norm = torch.load(ci_load_grad_norm_path)
|
|
assert math.isclose(
|
|
grad_norm,
|
|
expected_grad_norm,
|
|
rel_tol=0.01,
|
|
abs_tol=0.01,
|
|
), f"grad norm mismatch: {grad_norm} != {expected_grad_norm}"
|
|
# Close out pre-hooks if using distributed optimizer and overlapped param gather.
|
|
if pre_hook_enabled:
|
|
disable_forward_pre_hook(model)
|
|
|
|
|
|
def save(
|
|
iteration: int, model: Sequence[DDP], optimizer: MegatronOptimizer, opt_param_scheduler: OptimizerParamScheduler
|
|
) -> None:
|
|
"""Persist a training checkpoint safely with forward hooks disabled.
|
|
|
|
Args:
|
|
iteration (int): Current global iteration number.
|
|
model (Sequence[DDP]): Sequence of DDP-wrapped model chunks.
|
|
optimizer (MegatronOptimizer): Optimizer instance.
|
|
opt_param_scheduler (OptimizerParamScheduler): LR/WD scheduler.
|
|
"""
|
|
args = get_args()
|
|
if should_disable_forward_pre_hook(args):
|
|
disable_forward_pre_hook(model)
|
|
save_checkpoint(
|
|
iteration,
|
|
model,
|
|
optimizer,
|
|
opt_param_scheduler,
|
|
num_floating_point_operations_so_far=0,
|
|
checkpointing_context=None,
|
|
train_data_iterator=None,
|
|
preprocess_common_state_dict_fn=None,
|
|
)
|
|
if should_disable_forward_pre_hook(args):
|
|
enable_forward_pre_hook(model)
|
|
|
|
|
|
def initialize_model_and_optimizer(
|
|
args: Namespace, role: str = "actor"
|
|
) -> tuple[list[DDP], MegatronOptimizer, OptimizerParamScheduler, int]:
|
|
"""Initialize model(s), optimizer, scheduler, and load from checkpoint.
|
|
|
|
Args:
|
|
args (Namespace): Runtime arguments.
|
|
role (str): Logical role of the model (e.g., "actor", "critic").
|
|
|
|
Returns:
|
|
tuple[list[DDP], MegatronOptimizer, OptimizerParamScheduler, int]:
|
|
DDP-wrapped model chunks, optimizer, scheduler, and iteration index.
|
|
"""
|
|
|
|
if torch.version.hip:
|
|
import megatron.core.dist_checkpointing.strategies.filesystem_async as filesystem_async_module
|
|
from slime.utils.rocm_checkpoint_writer import ROCmFileSystemWriterAsync
|
|
|
|
filesystem_async_module.FileSystemWriterAsync = ROCmFileSystemWriterAsync
|
|
print("[ROCm] Applied FileSystemWriterAsync patch for HIP compatibility")
|
|
|
|
model, optimizer, opt_param_scheduler = setup_model_and_optimizer(args, role)
|
|
model[0].role = role
|
|
clear_memory()
|
|
iteration, _ = load_checkpoint(
|
|
model,
|
|
optimizer,
|
|
opt_param_scheduler,
|
|
checkpointing_context={},
|
|
skip_load_to_model_and_opt=False,
|
|
)
|
|
clear_memory()
|
|
|
|
opt_param_scheduler.step(increment=iteration * args.global_batch_size)
|
|
|
|
return model, optimizer, opt_param_scheduler, iteration
|