diff --git a/scripts/models/qwen36-35B-A3B.sh b/scripts/models/qwen36-35B-A3B.sh new file mode 100644 index 0000000000..b976dc3885 --- /dev/null +++ b/scripts/models/qwen36-35B-A3B.sh @@ -0,0 +1,4 @@ +# Qwen3.6-35B-A3B uses the same language/MoE topology as the Qwen3.5-35B-A3B +# config supported by Slime's qwen3_5 plugin. +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/qwen3.5-35B-A3B.sh" diff --git a/slime/backends/megatron_utils/__init__.py b/slime/backends/megatron_utils/__init__.py index b1936ae692..d4b22e1476 100644 --- a/slime/backends/megatron_utils/__init__.py +++ b/slime/backends/megatron_utils/__init__.py @@ -42,3 +42,6 @@ def _patched_forward(self, *args, packed_seq_params=None, **kwargs): logging.getLogger("megatron").setLevel(logging.WARNING) from . import megatron_patch # noqa: F401, E402 +from . import linear_ce_fallback # noqa: F401, E402 +from . import qwen_vl_packed_gdn # noqa: F401, E402 +from . import qwen_vl_text_language_fastpath # noqa: F401, E402 diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index f29afa0bd0..bcd3a3bf23 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -530,6 +530,9 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None: if self.args.debug_rollout_only: return + clear_memory(clear_host_memory=True) + print_memory("before save_model", clear_before_print=False) + # torch dist may trigger nccl communication during saving. if self.args.offload_train: self.wake_up() @@ -550,6 +553,9 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None: if self.args.offload_train: self.sleep() + clear_memory(clear_host_memory=True) + print_memory("after save_model", clear_before_print=False) + @timer def update_weights(self) -> None: if self.args.debug_train_only or self.args.debug_rollout_only: diff --git a/slime/backends/megatron_utils/arguments.py b/slime/backends/megatron_utils/arguments.py index 833af4a05a..22c3c328f8 100644 --- a/slime/backends/megatron_utils/arguments.py +++ b/slime/backends/megatron_utils/arguments.py @@ -3,9 +3,20 @@ from megatron.training.arguments import parse_args as _megatron_parse_args from megatron.training.arguments import validate_args as _megatron_validate_args -from megatron.training.tokenizer.tokenizer import _vocab_size_with_padding from transformers import AutoConfig +try: + from megatron.training.tokenizer.tokenizer import _vocab_size_with_padding +except ImportError: + + def _vocab_size_with_padding(orig_vocab_size, args): + after = int(orig_vocab_size) + multiple = int(getattr(args, "make_vocab_size_divisible_by", 128)) + multiple *= int(getattr(args, "tensor_model_parallel_size", 1)) + if multiple > 0: + after = ((after + multiple - 1) // multiple) * multiple + return after + __all__ = ["validate_args", "megatron_parse_args", "set_default_megatron_args"] logger = logging.getLogger(__name__) @@ -160,6 +171,8 @@ def _set_default_megatron_args(args): args.max_position_embeddings = args.seq_length # TODO: revisit this when megatron(dev) have solved the optimizer-cpu-offload ckpt saving bug args.dist_ckpt_save_pre_mcore_014 = True + if not hasattr(args, "enable_gloo_process_groups"): + args.enable_gloo_process_groups = True # compatible for megatron if hasattr(args, "rope_type") and args.rope_type is None: args.rope_type = "yarn" if args.multi_latent_attention else "rope" diff --git a/slime/backends/megatron_utils/cp_utils.py b/slime/backends/megatron_utils/cp_utils.py index 96c97df0e2..160cf3e26f 100644 --- a/slime/backends/megatron_utils/cp_utils.py +++ b/slime/backends/megatron_utils/cp_utils.py @@ -6,9 +6,26 @@ from megatron.core import mpu +def maybe_padded_total_lengths( + total_lengths: list[int], + qkv_format: str, + is_vl_model: bool, +) -> list[int] | None: + """Per-sample tp*cp*2 padded lengths for bridge VLM + CP + THD.""" + cp_size = mpu.get_context_parallel_world_size() + if not (is_vl_model and cp_size > 1 and qkv_format == "thd"): + return None + tp_size = mpu.get_tensor_model_parallel_world_size() + align = tp_size * cp_size * 2 + return [(t + align - 1) // align * align for t in total_lengths] + + def get_logits_and_tokens_offset_with_cp( total_length: int, response_length: int, + qkv_format: str = "thd", + max_seq_len: int | None = None, + padded_total_length: int | None = None, ): """ All offsets start from the begining of the prompt. @@ -18,7 +35,16 @@ def get_logits_and_tokens_offset_with_cp( assert cp_size > 1 prompt_length = total_length - response_length - chunk_size = (total_length + 2 * cp_size - 1) // (2 * cp_size) + if padded_total_length is not None: + assert padded_total_length % (2 * cp_size) == 0, ( + f"padded_total_length={padded_total_length} not divisible by 2*cp={2 * cp_size}" + ) + chunk_size = padded_total_length // (2 * cp_size) + elif qkv_format == "thd": + chunk_size = (total_length + 2 * cp_size - 1) // (2 * cp_size) + else: + assert max_seq_len is not None, "max_seq_len must be provided for qkv_format=bshd" + chunk_size = (max_seq_len + 2 * cp_size - 1) // (2 * cp_size) # the offset of 2 chunks chunk_0 = (cp_rank * chunk_size, (cp_rank + 1) * chunk_size) @@ -50,6 +76,9 @@ def get_sum_of_sample_mean( loss_masks: list[torch.Tensor], sample_denoms: list[torch.Tensor] | torch.Tensor | None = None, calculate_per_token_loss: bool = False, + qkv_format: str = "thd", + max_seq_lens: list[int] | None = None, + padded_total_lengths: list[int] | None = None, ) -> Callable[[torch.Tensor], torch.Tensor]: """ Calculate correct sample mean for CP. @@ -92,9 +121,15 @@ def sum_of_token(x: torch.Tensor) -> torch.Tensor: cp_chunk_lengths: list[int] = [] chunked_loss_masks: list[torch.Tensor] = [] - for total_length, response_length, loss_mask in zip(total_lengths, response_lengths, loss_masks, strict=False): + for i, (total_length, response_length, loss_mask) in enumerate( + zip(total_lengths, response_lengths, loss_masks, strict=False) + ): + max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None + padded_total_length = padded_total_lengths[i] if padded_total_lengths is not None else None prompt_length = total_length - response_length - _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp(total_length, response_length) + _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp( + total_length, response_length, qkv_format, max_seq_len, padded_total_length + ) loss_mask_0 = loss_mask[tokens_offset[0][0] - prompt_length : tokens_offset[0][1] - prompt_length] loss_mask_1 = loss_mask[tokens_offset[1][0] - prompt_length : tokens_offset[1][1] - prompt_length] chunked_loss_mask = torch.cat([loss_mask_0, loss_mask_1], dim=0) @@ -232,7 +267,12 @@ def gather_and_reduce_log_dict( return None -def all_gather_with_cp(tensor: torch.Tensor, total_length: int, response_length: int) -> torch.Tensor: +def all_gather_with_cp( + tensor: torch.Tensor, + total_length: int, + response_length: int, + padded_total_length: int | None = None, +) -> torch.Tensor: """ Gather tensors across all ranks in the context parallel group. The first dimension of the output tensor will be the `response_length`. @@ -243,7 +283,9 @@ def all_gather_with_cp(tensor: torch.Tensor, total_length: int, response_length: if cp_size == 1: return tensor - _, _, logits_offset, _ = get_logits_and_tokens_offset_with_cp(total_length, response_length) + _, _, logits_offset, _ = get_logits_and_tokens_offset_with_cp( + total_length, response_length, padded_total_length=padded_total_length + ) prompt_length = total_length - response_length @@ -287,10 +329,15 @@ def zero(len: int) -> torch.Tensor: def slice_with_cp( tokens: torch.Tensor, pad_value: tuple[int, float, Callable], + qkv_format: str = "thd", + max_seq_len: int | None = None, ) -> torch.Tensor: cp_rank = mpu.get_context_parallel_rank() cp_size = mpu.get_context_parallel_world_size() + if qkv_format == "bshd": + assert max_seq_len is not None + def pad_tokens(tokens, pad): if isinstance(pad_value, Callable): pad_func = pad_value @@ -302,10 +349,16 @@ def pad_tokens(tokens, pad): return tokens if cp_size == 1: + if qkv_format == "bshd": + pad = max_seq_len - tokens.size(0) + tokens = pad_tokens(tokens, pad) return tokens token_len = len(tokens) - chunk_size = (token_len + 2 * cp_size - 1) // (2 * cp_size) + if qkv_format == "thd": + chunk_size = (token_len + 2 * cp_size - 1) // (2 * cp_size) + else: + chunk_size = (max_seq_len + 2 * cp_size - 1) // (2 * cp_size) # pad pad = 2 * cp_size * chunk_size - token_len @@ -321,6 +374,9 @@ def slice_log_prob_with_cp( log_prob: list[float] | torch.Tensor, total_length: int, response_length: int, + qkv_format: str = "thd", + max_token_len: int | None = None, + padded_total_length: int | None = None, ) -> list[float] | torch.Tensor: assert len(log_prob) == response_length, ( f"log_prob length mismatch: len(log_prob)={len(log_prob)}, " @@ -333,7 +389,9 @@ def slice_log_prob_with_cp( return log_prob prompt_length = total_length - response_length - _, _, logits_offset, _ = get_logits_and_tokens_offset_with_cp(total_length, response_length) + _, _, logits_offset, _ = get_logits_and_tokens_offset_with_cp( + total_length, response_length, qkv_format, max_token_len, padded_total_length + ) chunk_1 = log_prob[logits_offset[0][0] - (prompt_length - 1) : logits_offset[0][1] - (prompt_length - 1)] chunk_2 = log_prob[logits_offset[1][0] - (prompt_length - 1) : logits_offset[1][1] - (prompt_length - 1)] diff --git a/slime/backends/megatron_utils/data.py b/slime/backends/megatron_utils/data.py index 00f319928f..aa257d5e0d 100644 --- a/slime/backends/megatron_utils/data.py +++ b/slime/backends/megatron_utils/data.py @@ -1,4 +1,5 @@ import logging +import os from argparse import Namespace from collections.abc import Sequence @@ -8,6 +9,8 @@ import torch.nn.functional as F from megatron.core import mpu from megatron.core.packed_seq_params import PackedSeqParams +from megatron.training.global_vars import get_args +from torch.nn.utils.rnn import pad_sequence from slime.utils import train_metric_utils from slime.utils.flops_utils import calculate_fwd_flops @@ -18,12 +21,78 @@ from .cp_utils import ( gather_and_reduce_log_dict, get_sum_of_sample_mean, + maybe_padded_total_lengths, rollout_log_metric_contribution, slice_with_cp, ) logger = logging.getLogger(__name__) +PAD_RULES = { + "input_features": dict( + transpose=(0, 2), + padding_value=0.0, + ), + "feature_attention_mask": dict( + transpose=(0, 1), + padding_value=0, + ), +} + + +def has_multimodal_train_inputs(multimodal_train_inputs) -> bool: + if multimodal_train_inputs is None: + return False + if isinstance(multimodal_train_inputs, dict): + return any(has_multimodal_train_inputs(value) for value in multimodal_train_inputs.values()) + if isinstance(multimodal_train_inputs, (list, tuple)): + return any(has_multimodal_train_inputs(value) for value in multimodal_train_inputs) + if torch.is_tensor(multimodal_train_inputs): + return multimodal_train_inputs.numel() > 0 + return True + + +def qwen_vl_unsplit_only_with_mm() -> bool: + return os.environ.get("SLIME_QWENVL_UNSPLIT_ONLY_WITH_MM", "0").lower() in {"1", "true", "yes", "on"} + + +def qwen_vl_text_language_fastpath() -> bool: + return os.environ.get("SLIME_QWENVL_TEXT_LANGUAGE_FASTPATH", "0").lower() in {"1", "true", "yes", "on"} + + +def qwen_vl_text_fastpath_requires_unsplit_input() -> bool: + args = get_args() + return ( + qwen_vl_text_language_fastpath() + and os.environ.get("SLIME_QWENVL_TEXT_FASTPATH_LOCAL_MROPE", "1").lower() not in {"1", "true", "yes", "on"} + and getattr(args, "position_embedding_type", "rope") == "mrope" + and getattr(args, "qkv_format", "thd") == "thd" + and mpu.get_context_parallel_world_size() > 1 + ) + + +def pad_and_flatten( + tensor_list, + transpose=None, + padding_value=0, +): + if len(tensor_list) == 1: + t = tensor_list[0] + return t, [t.size(0)] + + if transpose is not None: + tensor_list = [t.transpose(*transpose) for t in tensor_list] + + padded = pad_sequence(tensor_list, batch_first=True, padding_value=padding_value) + + if transpose is not None: + padded_list = [t.transpose(*transpose) for t in padded] + else: + padded_list = [t for t in padded] + + num_items = [t.size(0) for t in padded_list] + return torch.cat(padded_list, dim=0), num_items + def get_batch( data_iterator: "DataIterator", @@ -65,54 +134,118 @@ def get_batch( cp_size = mpu.get_context_parallel_world_size() cp_rank = mpu.get_context_parallel_rank() + qkv_format = getattr(get_args(), "qkv_format", "thd") + + if qkv_format == "bshd": + max_seqlen = batch["max_seq_lens"][0] + assert max([t.size(0) for t in tokens]) <= max_seqlen + + if cp_size > 1: + chunk_size = (max_seqlen + 2 * cp_size - 1) // (2 * cp_size) + padded_len = 2 * cp_size * chunk_size + unsplit = [F.pad(t, (0, padded_len - t.size(0)), value=pad_token_id) for t in tokens] + batch["unsplit_tokens"] = torch.stack(unsplit) + + tokens = [slice_with_cp(t, pad_token_id, qkv_format, max_seqlen) for t in tokens] + tokens = torch.stack(tokens) + packed_seq_params = None + + elif qkv_format == "thd": + is_vl_model = has_multimodal_train_inputs(batch.get("multimodal_train_inputs")) + uses_unsplit_forward = getattr(get_args(), "uses_unsplit_forward", False) + needs_unsplit_input = is_vl_model or ( + uses_unsplit_forward + and (not qwen_vl_unsplit_only_with_mm() or qwen_vl_text_fastpath_requires_unsplit_input()) + ) + if needs_unsplit_input and cp_size > 1: + tp_size = mpu.get_tensor_model_parallel_world_size() + align_size = tp_size * cp_size * 2 + device = tokens[0].device + + seqlens = torch.tensor([t.size(0) for t in tokens], dtype=torch.int32, device=device) + seqlens_padded = (seqlens + align_size - 1) // align_size * align_size + cu_seqlens_padded = torch.zeros(len(tokens) + 1, dtype=torch.int32, device=device) + cu_seqlens_padded[1:] = torch.cumsum(seqlens_padded, dim=0) + max_seqlen_padded = int(seqlens_padded.max().item()) + + unsplit_tokens = pad_sequence(tokens, batch_first=True, padding_value=pad_token_id) + unsplit_attention_mask = torch.zeros_like(unsplit_tokens, dtype=torch.bool) + for i, seqlen in enumerate(seqlens.tolist()): + unsplit_attention_mask[i, :seqlen] = True + + batch["unsplit_tokens"] = unsplit_tokens + batch["unsplit_attention_mask"] = unsplit_attention_mask + batch["vlm_packed_seq_params"] = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens_padded, + cu_seqlens_kv=cu_seqlens_padded, + max_seqlen_q=max_seqlen_padded, + max_seqlen_kv=max_seqlen_padded, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + ) + batch["padded_total_lengths"] = seqlens_padded.tolist() - if allgather_cp: - # DSA mode: concatenate all sequences first, then slice once with CP. - # We also pad the *global* concatenated stream to make per-rank chunks equal. - cu_seqlens_list: list[int] = [0] - for t in tokens: - cu_seqlens_list.append(cu_seqlens_list[-1] + t.size(0)) - - tokens = torch.cat(tokens, dim=0) - - # Pad global stream so (1) divisible by cp_size (equal chunks), - # (2) divisible by pad_size (reduce fragmentation). - global_pad_size = cp_size * pad_size - pad = (global_pad_size - tokens.size(0) % global_pad_size) % global_pad_size - if pad != 0: - tokens = F.pad(tokens, (0, pad), value=pad_token_id) - cu_seqlens_list.append(cu_seqlens_list[-1] + pad) - - cu_seqlens = torch.tensor(cu_seqlens_list, dtype=torch.int, device=torch.cuda.current_device()) - tokens = tokens.chunk(cp_size, dim=0)[cp_rank] + if allgather_cp: + # DSA mode: concatenate all sequences first, then slice once with CP. + # We also pad the *global* concatenated stream to make per-rank chunks equal. + cu_seqlens_list: list[int] = [0] + for t in tokens: + cu_seqlens_list.append(cu_seqlens_list[-1] + t.size(0)) + + tokens = torch.cat(tokens, dim=0) + global_pad_size = cp_size * pad_size + pad = (global_pad_size - tokens.size(0) % global_pad_size) % global_pad_size + if pad != 0: + tokens = F.pad(tokens, (0, pad), value=pad_token_id) + cu_seqlens_list.append(cu_seqlens_list[-1] + pad) + + cu_seqlens = torch.tensor(cu_seqlens_list, dtype=torch.int, device=torch.cuda.current_device()) + tokens = tokens.chunk(cp_size, dim=0)[cp_rank] + else: + tokens = [slice_with_cp(t, pad_token_id, qkv_format) for t in tokens] + + cu_seqlens = [0] + for t in tokens: + cu_seqlens.append(cu_seqlens[-1] + t.size(0)) + + tokens = torch.cat(tokens) + + # Always pad to reduce memory fragmentation and maybe make the computation faster + pad = (pad_size - tokens.size(0) % pad_size) % pad_size + if pad != 0: + tokens = F.pad(tokens, (0, pad), value=pad_token_id) + cu_seqlens.append(cu_seqlens[-1] + pad) + + # thd requires the cu_seqlens to be of the origin length + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int).cuda() * cp_size + + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() + packed_seq_params_kwargs = {} + if ( + cp_size > 1 + and uses_unsplit_forward + and qwen_vl_unsplit_only_with_mm() + and not is_vl_model + and not allgather_cp + ): + packed_seq_params_kwargs = { + "cu_seqlens_q_padded": cu_seqlens, + "cu_seqlens_kv_padded": cu_seqlens, + } + + packed_seq_params = PackedSeqParams( + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + qkv_format="thd", + **packed_seq_params_kwargs, + ) + + tokens = tokens.unsqueeze(0) else: - tokens = [slice_with_cp(t, pad_token_id) for t in tokens] - - cu_seqlens = [0] - for t in tokens: - cu_seqlens.append(cu_seqlens[-1] + t.size(0)) - - tokens = torch.cat(tokens) - - # Always pad to reduce memory fragmentation and maybe make the computation faster - pad = (pad_size - tokens.size(0) % pad_size) % pad_size - if pad != 0: - tokens = F.pad(tokens, (0, pad), value=pad_token_id) - cu_seqlens.append(cu_seqlens[-1] + pad) - - # thd requires the cu_seqlens to be of the origin length - cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int).cuda() * cp_size - - max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() - packed_seq_params = PackedSeqParams( - cu_seqlens_q=cu_seqlens, - cu_seqlens_kv=cu_seqlens, - max_seqlen_q=max_seqlen, - max_seqlen_kv=max_seqlen, - qkv_format="thd", - ) - - tokens = tokens.unsqueeze(0) + raise ValueError(f"Unsupported qkv_format: {qkv_format}") batch["tokens"] = tokens batch["packed_seq_params"] = packed_seq_params @@ -131,16 +264,18 @@ def get_batch( if allgather_cp: loss_masks.append(loss_mask) continue - loss_mask = slice_with_cp(loss_mask, 0) + loss_mask = slice_with_cp(loss_mask, 0, qkv_format, max_seqlen) loss_masks.append(loss_mask) - if allgather_cp: + if qkv_format == "bshd": + loss_masks = torch.stack(loss_masks) + elif qkv_format == "thd" and allgather_cp: # DSA: concatenate first (same as tokens), pad globally (same pad as above), then slice once. loss_masks = torch.cat(loss_masks, dim=0) if pad != 0: loss_masks = F.pad(loss_masks, (0, pad), value=0) loss_masks = loss_masks.chunk(cp_size, dim=0)[cp_rank].unsqueeze(0) - else: + elif qkv_format == "thd": loss_masks = torch.cat(loss_masks) loss_masks = F.pad(loss_masks, (0, pad), value=0).unsqueeze(0) @@ -149,18 +284,37 @@ def get_batch( # Process multimodal training tensors if present multimodal_train_inputs = batch.get("multimodal_train_inputs", None) - if multimodal_train_inputs is not None: + if has_multimodal_train_inputs(multimodal_train_inputs): multimodal_data = {} # key -> concatenated tensor + multimodal_num_items = {} + tensor_dict_list = {} for mm_input_dict in multimodal_train_inputs: - if mm_input_dict is not None: - for key, mm_tensor in mm_input_dict.items(): - if key not in multimodal_data: - multimodal_data[key] = mm_tensor - else: - multimodal_data[key] = torch.cat([multimodal_data[key], mm_tensor], dim=0) - batch["multimodal_train_inputs"] = multimodal_data + if mm_input_dict is None: + continue + for key, mm_tensor in mm_input_dict.items(): + if isinstance(mm_tensor, list): + mm_tensor = torch.tensor(mm_tensor, device=batch["tokens"].device) + tensor_dict_list.setdefault(key, []).append(mm_tensor) + + for key, tensor_list in tensor_dict_list.items(): + if key in PAD_RULES: + multimodal_data[key], multimodal_num_items[key] = pad_and_flatten( + tensor_list, + **PAD_RULES[key], + ) + else: + if len(tensor_list) == 1: + multimodal_data[key] = tensor_list[0] + else: + multimodal_data[key] = torch.cat(tensor_list, dim=0) + multimodal_num_items[key] = [t.size(0) for t in tensor_list] + batch["multimodal_train_inputs"] = multimodal_data or None + batch["multimodal_num_items"] = multimodal_num_items + else: + batch["multimodal_train_inputs"] = None + batch["multimodal_num_items"] = {} - return batch + return move_tensors_to_device(batch, batch["tokens"].device) def gather_log_data( @@ -238,6 +392,18 @@ def reset(self) -> "DataIterator": return self +def move_tensors_to_device(data, device): + if isinstance(data, torch.Tensor): + return data.to(device) + if isinstance(data, dict): + return {k: move_tensors_to_device(v, device) for k, v in data.items()} + if isinstance(data, list): + return [move_tensors_to_device(v, device) for v in data] + if isinstance(data, tuple): + return tuple(move_tensors_to_device(v, device) for v in data) + return data + + def get_data_iterator(rollout_data: RolloutBatch) -> list[DataIterator]: """Build one ``DataIterator`` per VPP stage from the pre-computed schedule in ``rollout_data``.""" vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1 @@ -265,6 +431,12 @@ def log_rollout_data( response_lengths = rollout_data["response_lengths"] loss_masks = rollout_data["loss_masks"] total_lengths = rollout_data["total_lengths"] + max_seq_lens = rollout_data.get("max_seq_lens", None) + padded_total_lengths = maybe_padded_total_lengths( + total_lengths, + args.qkv_format, + rollout_data.get("multimodal_train_inputs") is not None or getattr(args, "uses_unsplit_forward", False), + ) # Same per-rollout denominators the training loss uses, so reported # log_probs / returns / advantages / etc. live in the same per-rollout # mean space (rather than per-sample) as the gradient signal. @@ -316,6 +488,9 @@ def log_rollout_data( response_lengths, loss_masks, rollout_mask_sums, + qkv_format=args.qkv_format, + max_seq_lens=max_seq_lens, + padded_total_lengths=padded_total_lengths, ) # Compute (sum, count) via the shared helper so this # path and the unit tests stay in sync. @@ -399,15 +574,24 @@ def quantile(total_value, n_quantiles, data) -> dict: raw_rewards = rollout_data["raw_reward"] # Additional metrics for correct cases are calculated separately below. + padded_total_lengths = maybe_padded_total_lengths( + total_lengths, + args.qkv_format, + rollout_data.get("multimodal_train_inputs") is not None + or getattr(args, "uses_unsplit_forward", False), + ) correct_response_lengths = [] correct_total_lengths = [] correct_loss_masks = [] + correct_padded_total_lengths = [] if padded_total_lengths is not None else None correct_entropy = [] for i, raw_reward in enumerate(raw_rewards): if raw_reward == 1: correct_response_lengths.append(response_lengths[i]) correct_total_lengths.append(total_lengths[i]) correct_loss_masks.append(loss_masks[i]) + if correct_padded_total_lengths is not None: + correct_padded_total_lengths.append(padded_total_lengths[i]) correct_entropy.append(-rollout_data["log_probs"][i]) num_correct_responses = len(correct_total_lengths) rollout_data["correct_response_lengths"] = correct_response_lengths @@ -424,7 +608,12 @@ def quantile(total_value, n_quantiles, data) -> dict: # entropy report. Per-sample-mean over the filtered subset is # the cleanest semantic. sum_of_sample_mean = get_sum_of_sample_mean( - correct_total_lengths, correct_response_lengths, correct_loss_masks, sample_denoms=None + correct_total_lengths, + correct_response_lengths, + correct_loss_masks, + sample_denoms=None, + qkv_format=args.qkv_format, + padded_total_lengths=correct_padded_total_lengths, ) correct_entropy = sum_of_sample_mean(torch.cat(correct_entropy, dim=0)) rollout_data["correct_entropy"] = [correct_entropy.item()] * num_correct_responses diff --git a/slime/backends/megatron_utils/initialize.py b/slime/backends/megatron_utils/initialize.py index 36872bdd20..344a619da4 100644 --- a/slime/backends/megatron_utils/initialize.py +++ b/slime/backends/megatron_utils/initialize.py @@ -11,6 +11,12 @@ logger = logging.getLogger(__name__) +def get_use_gloo_process_groups(args) -> bool: + if hasattr(args, "enable_gloo_process_groups"): + return args.enable_gloo_process_groups + return getattr(args, "use_gloo_process_groups", False) + + def _set_random_seed( seed_: int, data_parallel_random_init: bool = False, @@ -49,7 +55,7 @@ def _initialize_distributed(args, get_embedding_ranks=None, get_position_embeddi order="tp-cp-ep-dp-pp" if not args.use_tp_pp_dp_mapping else "tp-cp-ep-pp-dp", get_embedding_ranks=get_embedding_ranks, get_position_embedding_ranks=get_position_embedding_ranks, - create_gloo_process_groups=args.enable_gloo_process_groups, + create_gloo_process_groups=get_use_gloo_process_groups(args), ) diff --git a/slime/backends/megatron_utils/linear_ce_fallback.py b/slime/backends/megatron_utils/linear_ce_fallback.py new file mode 100644 index 0000000000..84f329df2d --- /dev/null +++ b/slime/backends/megatron_utils/linear_ce_fallback.py @@ -0,0 +1,434 @@ +"""Runtime fallback for Megatron linear cross entropy. + +Slime v0.3.0 can run on images where Megatron's Blackwell linear CE module is +present but its CUDA/Cutlass entry points were not built. Megatron then fails +late at the first long-context SFT forward. This patch keeps the linear-CE +code path enabled and provides a chunked PyTorch implementation only when the +native entry points are absent. +""" + +from __future__ import annotations + +import logging +import os +import time +from typing import Literal + +import torch +import torch.distributed as dist + +logger = logging.getLogger(__name__) +_SHAPE_MISMATCH_LOGGED = False +_TIMING_COUNTS = {"forward": 0, "backward": 0} +_TIMING_CONFIG_LOGGED = False + + +def _env_value(name: str, default: str) -> str: + return os.environ.get(name, os.environ.get(name.lower(), default)) + + +def _env_flag(name: str) -> bool: + return _env_value(name, "0").lower() in {"1", "true", "yes", "on"} + + +def _env_int(name: str, default: int) -> int: + try: + return int(_env_value(name, str(default))) + except ValueError: + return default + + +def _rank() -> int: + try: + if dist.is_available() and dist.is_initialized(): + return dist.get_rank() + except Exception: + pass + return int(os.environ.get("RANK", "0")) + + +def _timing_should_sample(phase: Literal["forward", "backward"]) -> tuple[bool, int]: + global _TIMING_CONFIG_LOGGED + if not _env_flag("SLIME_LINEAR_CE_FALLBACK_TIMING"): + return False, 0 + rank_filter = _env_int("SLIME_LINEAR_CE_FALLBACK_TIMING_RANK", 0) + rank = _rank() + if rank_filter >= 0 and rank != rank_filter: + return False, 0 + + if not _TIMING_CONFIG_LOGGED: + _TIMING_CONFIG_LOGGED = True + logger.info( + "Linear CE fallback timing enabled: rank=%s interval=%s limit=%s vocab_chunk=%s", + rank_filter, + _env_int("SLIME_LINEAR_CE_FALLBACK_TIMING_INTERVAL", 32), + _env_int("SLIME_LINEAR_CE_FALLBACK_TIMING_LIMIT", 64), + _chunk_size(), + ) + + count = _TIMING_COUNTS.get(phase, 0) + 1 + _TIMING_COUNTS[phase] = count + limit = _env_int("SLIME_LINEAR_CE_FALLBACK_TIMING_LIMIT", 64) + interval = max(1, _env_int("SLIME_LINEAR_CE_FALLBACK_TIMING_INTERVAL", 32)) + return count <= limit and (count == 1 or count % interval == 0), count + + +def _timing_start(enabled: bool): + if not enabled: + return None, 0.0 + if torch.cuda.is_available(): + event = torch.cuda.Event(enable_timing=True) + event.record() + return event, 0.0 + return None, time.perf_counter() + + +def _timing_elapsed_ms(start_event, start_wall: float) -> float: + if start_event is not None: + end_event = torch.cuda.Event(enable_timing=True) + end_event.record() + torch.cuda.synchronize() + return float(start_event.elapsed_time(end_event)) + return (time.perf_counter() - start_wall) * 1000.0 + + +def _world_info(tp_group): + if tp_group is None: + return 0, 1 + return dist.get_rank(tp_group), dist.get_world_size(tp_group) + + +def _chunk_size() -> int: + return max(1, int(os.environ.get("SLIME_LINEAR_CE_FALLBACK_VOCAB_CHUNK", "1024"))) + + +def _labels_for_local_vocab(labels: torch.Tensor, vocab_size: int, tp_rank: int) -> torch.Tensor: + return labels - tp_rank * vocab_size + + +def _valid_labels(labels: torch.Tensor, ignore_index: int, vocab_size: int, tp_world_size: int) -> torch.Tensor: + global_vocab_size = vocab_size * tp_world_size + return labels.ne(ignore_index) & labels.ge(0) & labels.lt(global_vocab_size) + + +def _assert_no_valid_tail_labels( + labels: torch.Tensor, + *, + common_tokens: int, + ignore_index: int, + vocab_size: int, + tp_world_size: int, + hidden_tokens: int, +) -> None: + if labels.numel() <= common_tokens: + return + tail = labels[common_tokens:] + valid_tail = _valid_labels(tail, ignore_index, vocab_size, tp_world_size) + if valid_tail.any(): + first = torch.nonzero(valid_tail, as_tuple=False).flatten()[0].item() + raise RuntimeError( + "Linear CE fallback would drop valid tail labels because hidden is shorter than labels: " + f"hidden_tokens={hidden_tokens}, label_tokens={labels.numel()}, common_tokens={common_tokens}, " + f"first_valid_tail_offset={common_tokens + first}, label={tail[first].item()}" + ) + + +def _reduce_loss(losses: torch.Tensor, valid: torch.Tensor, reduction: Literal["none", "sum", "mean"]): + losses = losses.masked_fill(~valid, 0.0) + if reduction == "none": + return losses + if reduction == "sum": + return losses.sum() + if reduction == "mean": + denom = torch.clamp_min(valid.sum(), 1) + return losses.sum() / denom + raise ValueError(f"Unsupported reduction: {reduction}") + + +def _forward( + hidden: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + tp_group=None, + reduction: Literal["none", "sum", "mean"] = "mean", + ignore_index: int = -100, + sequence_parallel: bool = False, +): + should_time, timing_sample = _timing_should_sample("forward") + timing_event, timing_wall = _timing_start(should_time) + + global _SHAPE_MISMATCH_LOGGED + tp_rank, tp_world_size = _world_info(tp_group) + in_tp_mode = tp_group is not None and tp_world_size > 1 + + hidden_shape = hidden.shape + hidden_view = hidden.view(-1, hidden.shape[-1]) + labels_view = labels.view(-1) + global_hidden = hidden + + if in_tp_mode and sequence_parallel: + partial_shape = hidden.shape + global_shape = (partial_shape[0] * tp_world_size, *partial_shape[1:]) + global_hidden = torch.empty(global_shape, dtype=hidden.dtype, device=hidden.device) + dist.all_gather_into_tensor(global_hidden, hidden, group=tp_group) + hidden_view = global_hidden.view(-1, global_hidden.shape[-1]) + + label_tokens = labels_view.numel() + hidden_tokens = hidden_view.shape[0] + if hidden_tokens != label_tokens and not _SHAPE_MISMATCH_LOGGED: + logger.warning( + "Linear CE fallback aligning padded hidden/label rows: hidden_tokens=%s labels=%s " + "hidden_shape=%s labels_shape=%s sequence_parallel=%s tp_world_size=%s", + hidden_tokens, + label_tokens, + tuple(hidden_shape), + tuple(labels.shape), + sequence_parallel, + tp_world_size, + ) + _SHAPE_MISMATCH_LOGGED = True + common_tokens = min(hidden_tokens, label_tokens) + hidden_used = hidden_view[:common_tokens] + labels_used = labels_view[:common_tokens] + + num_tokens = common_tokens + vocab_size = weight.shape[0] + chunk = _chunk_size() + _assert_no_valid_tail_labels( + labels_view, + common_tokens=common_tokens, + ignore_index=ignore_index, + vocab_size=vocab_size, + tp_world_size=tp_world_size, + hidden_tokens=hidden_tokens, + ) + labels_local = _labels_for_local_vocab(labels_used, vocab_size, tp_rank) + valid = _valid_labels(labels_used, ignore_index, vocab_size, tp_world_size) + + local_max = torch.full((num_tokens,), -float("inf"), dtype=torch.float32, device=hidden.device) + for start in range(0, vocab_size, chunk): + weight_chunk = weight[start : start + chunk] + logits = torch.matmul(hidden_used, weight_chunk.t()).float() + local_max = torch.maximum(local_max, logits.max(dim=-1).values) + + maximum = local_max + if in_tp_mode: + maximum = maximum.clone() + dist.all_reduce(maximum, op=dist.ReduceOp.MAX, group=tp_group) + + accumulate = torch.zeros((num_tokens,), dtype=torch.float32, device=hidden.device) + target_logits = torch.zeros((num_tokens,), dtype=torch.float32, device=hidden.device) + for start in range(0, vocab_size, chunk): + end = min(start + chunk, vocab_size) + weight_chunk = weight[start:end] + logits = torch.matmul(hidden_used, weight_chunk.t()).float() + accumulate.add_(torch.exp(logits - maximum[:, None]).sum(dim=-1)) + + in_chunk = valid & labels_local.ge(start) & labels_local.lt(end) + if in_chunk.any(): + row_idx = torch.nonzero(in_chunk, as_tuple=False).flatten() + col_idx = labels_local[row_idx] - start + target_logits[row_idx] = logits[row_idx, col_idx].float() + + if in_tp_mode: + dist.all_reduce(accumulate, op=dist.ReduceOp.SUM, group=tp_group) + dist.all_reduce(target_logits, op=dist.ReduceOp.SUM, group=tp_group) + + losses = torch.log(accumulate) + maximum - target_logits + if reduction == "none" and label_tokens > common_tokens: + padded_losses = torch.zeros((label_tokens,), dtype=losses.dtype, device=losses.device) + padded_losses[:common_tokens] = losses + padded_valid = torch.zeros((label_tokens,), dtype=torch.bool, device=valid.device) + padded_valid[:common_tokens] = valid + logprobs = _reduce_loss(padded_losses, padded_valid, reduction) + else: + logprobs = _reduce_loss(losses, valid, reduction) + num_valid_tokens = valid.sum().to(dtype=torch.int64) + + if should_time: + elapsed_ms = _timing_elapsed_ms(timing_event, timing_wall) + logger.info( + "Linear CE fallback timing: phase=forward sample=%s rank=%s elapsed_ms=%.3f " + "hidden_shape=%s weight_shape=%s labels_shape=%s common_tokens=%s valid_tokens=%s " + "vocab_size=%s chunk=%s reduction=%s sequence_parallel=%s tp_world_size=%s", + timing_sample, + _rank(), + elapsed_ms, + tuple(hidden_shape), + tuple(weight.shape), + tuple(labels.shape), + common_tokens, + int(num_valid_tokens.detach().item()), + vocab_size, + chunk, + reduction, + sequence_parallel, + tp_world_size, + ) + + return ( + logprobs, + maximum, + accumulate, + num_valid_tokens, + tp_rank, + tp_world_size, + global_hidden, + ) + + +def _grad_factor( + dlogprobs: torch.Tensor, + labels_view: torch.Tensor, + valid: torch.Tensor, + num_valid_tokens: torch.Tensor, + reduction: Literal["none", "sum", "mean"], +) -> torch.Tensor: + if reduction == "none": + grad = dlogprobs.view(-1).float() + elif reduction == "sum": + grad = torch.ones_like(labels_view, dtype=torch.float32) * dlogprobs.float() + elif reduction == "mean": + denom = torch.clamp_min(num_valid_tokens, 1).float() + grad = torch.ones_like(labels_view, dtype=torch.float32) * dlogprobs.float() / denom + else: + raise ValueError(f"Unsupported reduction: {reduction}") + return grad.masked_fill(~valid, 0.0) + + +def _backward( + dlogprobs: torch.Tensor, + global_hidden: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + maximum: torch.Tensor, + accumulate: torch.Tensor, + num_valid_tokens: torch.Tensor, + reduction: Literal["none", "sum", "mean"] = "mean", + ignore_index: int = -100, + tp_group=None, + tp_rank: int = 0, + tp_world_size: int = 1, + sequence_parallel: bool = False, +): + should_time, timing_sample = _timing_should_sample("backward") + timing_event, timing_wall = _timing_start(should_time) + + in_tp_mode = tp_group is not None and tp_world_size > 1 + hidden_view = global_hidden.view(-1, global_hidden.shape[-1]) + labels_view = labels.view(-1) + label_tokens = labels_view.numel() + hidden_tokens = hidden_view.shape[0] + common_tokens = min(hidden_tokens, label_tokens) + hidden_used = hidden_view[:common_tokens] + labels_used = labels_view[:common_tokens] + vocab_size = weight.shape[0] + chunk = _chunk_size() + _assert_no_valid_tail_labels( + labels_view, + common_tokens=common_tokens, + ignore_index=ignore_index, + vocab_size=vocab_size, + tp_world_size=tp_world_size, + hidden_tokens=hidden_tokens, + ) + labels_local = _labels_for_local_vocab(labels_used, vocab_size, tp_rank) + valid = _valid_labels(labels_used, ignore_index, vocab_size, tp_world_size) + dlogprobs_used = dlogprobs.view(-1)[:common_tokens] if reduction == "none" else dlogprobs + grad = _grad_factor(dlogprobs_used, labels_used, valid, num_valid_tokens, reduction) + + d_hidden_used = torch.zeros_like(hidden_used) + d_weight = torch.zeros_like(weight) + + for start in range(0, vocab_size, chunk): + end = min(start + chunk, vocab_size) + weight_chunk = weight[start:end] + logits = torch.matmul(hidden_used, weight_chunk.t()).float() + probs = torch.exp(logits - maximum[:, None]) / accumulate[:, None] + + in_chunk = valid & labels_local.ge(start) & labels_local.lt(end) + if in_chunk.any(): + row_idx = torch.nonzero(in_chunk, as_tuple=False).flatten() + col_idx = labels_local[row_idx] - start + probs[row_idx, col_idx] -= 1.0 + + probs.mul_(grad[:, None]) + probs_for_mm = probs.to(dtype=hidden_used.dtype) + d_hidden_used.addmm_(probs_for_mm, weight_chunk) + d_weight[start:end] = torch.matmul(probs_for_mm.t(), hidden_used) + + d_hidden = torch.zeros_like(hidden_view) + d_hidden[:common_tokens] = d_hidden_used + + if in_tp_mode: + dist.all_reduce(d_hidden, op=dist.ReduceOp.SUM, group=tp_group) + if sequence_parallel: + partial_hidden_shape = (global_hidden.shape[0] // tp_world_size, *global_hidden.shape[1:]) + local_tokens = d_hidden.shape[0] // tp_world_size + d_hidden = d_hidden[tp_rank * local_tokens : (tp_rank + 1) * local_tokens].clone() + d_hidden = d_hidden.view(partial_hidden_shape) + + if should_time: + elapsed_ms = _timing_elapsed_ms(timing_event, timing_wall) + logger.info( + "Linear CE fallback timing: phase=backward sample=%s rank=%s elapsed_ms=%.3f " + "dlogprobs_shape=%s hidden_shape=%s weight_shape=%s labels_shape=%s common_tokens=%s " + "vocab_size=%s chunk=%s reduction=%s sequence_parallel=%s tp_world_size=%s", + timing_sample, + _rank(), + elapsed_ms, + tuple(dlogprobs.shape), + tuple(global_hidden.shape), + tuple(weight.shape), + tuple(labels.shape), + common_tokens, + vocab_size, + chunk, + reduction, + sequence_parallel, + tp_world_size, + ) + + return d_hidden.view_as(global_hidden if not (in_tp_mode and sequence_parallel) else d_hidden), d_weight + + +def apply_linear_ce_fallback_patch() -> None: + try: + from megatron.core.fusions import fused_linear_cross_entropy as fused_lce + except Exception: + return + + entry = None + try: + from megatron.core.fusions.linear_cross_entropy.blackwell import entry as blackwell_entry + + entry = blackwell_entry + except Exception: + pass + + missing_forward = entry is None or not callable(getattr(entry, "forward", None)) + missing_backward = entry is None or not callable(getattr(entry, "backward", None)) + if not (missing_forward or missing_backward): + return + + class FallbackPlatform: + def __init__(self) -> None: + self.forward_func = _forward + self.backward_func = _backward + + if entry is not None: + entry.forward = _forward + entry.backward = _backward + fused_lce.Platform = FallbackPlatform + try: + fused_lce._get_platform.cache_clear() + except Exception: + pass + logger.warning( + "Megatron linear_cross_entropy Blackwell entry points are unavailable; " + "using Slime chunked PyTorch fallback with vocab_chunk=%s.", + _chunk_size(), + ) + + +apply_linear_ce_fallback_patch() diff --git a/slime/backends/megatron_utils/loss.py b/slime/backends/megatron_utils/loss.py index adfa5d67e6..ae14e68a2a 100644 --- a/slime/backends/megatron_utils/loss.py +++ b/slime/backends/megatron_utils/loss.py @@ -1,3 +1,4 @@ +import os from argparse import Namespace from collections.abc import Callable, Iterator from typing import Any @@ -12,6 +13,7 @@ from slime.utils.misc import load_function from slime.utils.ppo_utils import ( calculate_log_probs_and_entropy, + calculate_log_probs_and_entropy_fused, compute_approx_kl, compute_cispo_loss, compute_gspo_kl, @@ -51,6 +53,17 @@ def get_rollout_top_p_logprob_kwargs(args: Namespace, batch: dict[str, Any]) -> } +_DEFERRED_LM_HEAD = {"weight": None} + + +def set_deferred_lm_head_weight(weight) -> None: + _DEFERRED_LM_HEAD["weight"] = weight + + +def _chunked_lm_head_enabled() -> bool: + return os.environ.get("CHUNKED_LM_HEAD", "0") == "1" + + def get_responses( logits: torch.Tensor, *, @@ -58,6 +71,8 @@ def get_responses( unconcat_tokens: list[torch.Tensor], total_lengths: list[int], response_lengths: list[int], + max_seq_lens: list[int] | None = None, + padded_total_lengths: list[int] | None = None, apply_temperature: bool = True, ) -> Iterator[tuple[torch.Tensor, torch.Tensor]]: """Yield response-aligned `(logits_chunk, tokens_chunk)` pairs per sample. @@ -82,10 +97,16 @@ def get_responses( `[R, V]` (policy) or `[R, 1]` (value) and `tokens_chunk` is shape `[R]` (1D int64), both aligned to response tokens for one sample. """ + qkv_format = args.qkv_format + assert logits.dtype == torch.float32, f"{logits.dtype}" assert len(logits.shape) == 3, f"{logits.shape}" - assert logits.size(0) == 1, f"{logits.shape}" - logits = logits.squeeze(0) + if qkv_format == "thd": + assert logits.size(0) == 1, f"{logits.shape}" + logits = logits.squeeze(0) + else: + assert max_seq_lens is not None + logits = logits.view(-1, logits.size(-1)) if apply_temperature and args.rollout_temperature != 1.0: logits = logits.div(args.rollout_temperature) @@ -93,10 +114,19 @@ def get_responses( cp_size = mpu.get_context_parallel_world_size() end = 0 seq_start = 0 - for tokens, total_length, response_length in zip(unconcat_tokens, total_lengths, response_lengths, strict=False): + for i, (tokens, total_length, response_length) in enumerate( + zip(unconcat_tokens, total_lengths, response_lengths, strict=False) + ): + max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None + padded_total_length = padded_total_lengths[i] if padded_total_lengths is not None else None + if cp_size == 1: - end += total_length - start = end - response_length + if qkv_format == "bshd": + end = max_seq_len * i + total_length + start = end - response_length + else: + end += total_length + start = end - response_length logits_chunk = logits[start - 1 : end - 1] tokens_chunk = tokens[-response_length:] elif args.allgather_cp: @@ -125,7 +155,7 @@ def get_responses( else: # TODO: this is super ugly... do better abstraction. chunk_size, chunks_offset, logits_offset, tokens_offset = get_logits_and_tokens_offset_with_cp( - total_length, response_length + total_length, response_length, qkv_format, max_seq_len, padded_total_length ) logits_0, logits_1 = logits[end : end + chunk_size], logits[end + chunk_size : end + 2 * chunk_size] @@ -152,8 +182,11 @@ def _allgather_cp_redistribute( res: dict[str, list[torch.Tensor]], *, logits_local_len: int, + args: Namespace, total_lengths: list[int], response_lengths: list[int], + max_seq_lens: list[int] | None = None, + padded_total_lengths: list[int] | None = None, ) -> None: """Redistribute response tensors from allgather-CP layout to zigzag ring-attn layout. @@ -219,10 +252,21 @@ def _allgather_cp_redistribute( # Re-slice each sample into zigzag CP pattern new_values = [] - for full_resp, total_length, response_length in zip( - all_cat.split(response_lengths, dim=0), total_lengths, response_lengths, strict=False + for idx, (full_resp, total_length, response_length) in enumerate( + zip(all_cat.split(response_lengths, dim=0), total_lengths, response_lengths, strict=False) ): - new_values.append(slice_log_prob_with_cp(full_resp, total_length, response_length)) + max_seq_len = max_seq_lens[idx] if max_seq_lens is not None else None + padded_total_length = padded_total_lengths[idx] if padded_total_lengths is not None else None + new_values.append( + slice_log_prob_with_cp( + full_resp, + total_length, + response_length, + args.qkv_format, + max_seq_len, + padded_total_length, + ) + ) res[key] = new_values @@ -233,7 +277,10 @@ def _build_shifted_tokens( unconcat_tokens: list[torch.Tensor], total_lengths: list[int], response_lengths: list[int], + qkv_format: str, + max_seq_lens: list[int] | None, allgather_cp: bool, + padded_total_lengths: list[int] | None = None, ) -> torch.Tensor: """Build shifted target tokens for the full packed/padded logits.""" cp_size = mpu.get_context_parallel_world_size() @@ -242,11 +289,13 @@ def _build_shifted_tokens( if cp_size > 1 and not allgather_cp: full_tokens = torch.zeros(T, dtype=torch.long, device=device) end = 0 - for tokens, total_length, response_length in zip( - unconcat_tokens, total_lengths, response_lengths, strict=False + for i, (tokens, total_length, response_length) in enumerate( + zip(unconcat_tokens, total_lengths, response_lengths, strict=False) ): + max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None + padded_total_length = padded_total_lengths[i] if padded_total_lengths is not None else None chunk_size_cp, chunks_offset, logits_offset, tokens_offset = get_logits_and_tokens_offset_with_cp( - total_length, response_length + total_length, response_length, qkv_format, max_seq_len, padded_total_length ) for half, base in ((0, end), (1, end + chunk_size_cp)): lo = logits_offset[half][0] - chunks_offset[half][0] @@ -259,10 +308,16 @@ def _build_shifted_tokens( T_global = sum(total_lengths) if allgather_cp else T full_tokens = torch.zeros(T_global, dtype=torch.long, device=device) - offset = 0 - for tokens, total_length in zip(unconcat_tokens, total_lengths, strict=False): - full_tokens[offset : offset + total_length - 1] = tokens[1:total_length] - offset += total_length + if qkv_format == "thd" or allgather_cp: + offset = 0 + for tokens, total_length in zip(unconcat_tokens, total_lengths, strict=False): + full_tokens[offset : offset + total_length - 1] = tokens[1:total_length] + offset += total_length + else: + assert max_seq_lens is not None + for i, (tokens, total_length) in enumerate(zip(unconcat_tokens, total_lengths, strict=False)): + seq_start = max_seq_lens[i] * i + full_tokens[seq_start : seq_start + total_length - 1] = tokens[1:total_length] # allgather-CP: slice to local chunk if allgather_cp: @@ -391,7 +446,10 @@ def _extract_per_sample( entropy_full: torch.Tensor | None, total_lengths: list[int], response_lengths: list[int], + qkv_format: str, + max_seq_lens: list[int] | None, allgather_cp: bool, + padded_total_lengths: list[int] | None = None, ) -> tuple[list[torch.Tensor], list[torch.Tensor | None]]: """Slice per-sample response log-probs/entropy from full-length 1-D tensors.""" cp_size = mpu.get_context_parallel_world_size() @@ -401,9 +459,11 @@ def _extract_per_sample( if cp_size > 1 and not allgather_cp: # zigzag CP pos = 0 - for total_length, response_length in zip(total_lengths, response_lengths, strict=False): + for i, (total_length, response_length) in enumerate(zip(total_lengths, response_lengths, strict=False)): + max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None + padded_total_length = padded_total_lengths[i] if padded_total_lengths is not None else None chunk_size_cp, chunks_offset, logits_offset, _tokens_offset = get_logits_and_tokens_offset_with_cp( - total_length, response_length + total_length, response_length, qkv_format, max_seq_len, padded_total_length ) lo0 = logits_offset[0][0] - chunks_offset[0][0] hi0 = logits_offset[0][1] - chunks_offset[0][0] @@ -455,14 +515,23 @@ def _extract_per_sample( else: # cp1 - offset = 0 - for total_length, response_length in zip(total_lengths, response_lengths, strict=False): - end = offset + total_length - start = end - response_length - log_probs_list.append(log_prob_full[start - 1 : end - 1]) - if entropy_full is not None: - entropy_list.append(entropy_full[start - 1 : end - 1]) - offset += total_length + if qkv_format == "thd": + offset = 0 + for total_length, response_length in zip(total_lengths, response_lengths, strict=False): + end = offset + total_length + start = end - response_length + log_probs_list.append(log_prob_full[start - 1 : end - 1]) + if entropy_full is not None: + entropy_list.append(entropy_full[start - 1 : end - 1]) + offset += total_length + else: + assert max_seq_lens is not None + for i, (total_length, response_length) in enumerate(zip(total_lengths, response_lengths, strict=False)): + end = max_seq_lens[i] * i + total_length + start = end - response_length + log_probs_list.append(log_prob_full[start - 1 : end - 1]) + if entropy_full is not None: + entropy_list.append(entropy_full[start - 1 : end - 1]) return log_probs_list, entropy_list @@ -476,8 +545,9 @@ def get_log_probs_and_entropy( response_lengths: list[int], with_entropy: bool = False, non_loss_data: bool = True, - top_p_token_ids: list[list[int]] | None = None, - top_p_token_offsets: list[list[int]] | None = None, + max_seq_lens: list[int] | None = None, + padded_total_lengths: list[int] | None = None, + lm_head_weight: torch.Tensor | None = None, ) -> dict[str, list[torch.Tensor]]: """Compute per-token log-probabilities (and optionally entropy) on responses. @@ -489,7 +559,9 @@ def get_log_probs_and_entropy( log-probabilities; entropy is always computed from the unmasked logits. """ assert non_loss_data - assert logits.dtype == torch.float32, f"{logits.dtype}" + qkv_format = args.qkv_format + + assert lm_head_weight is not None or logits.dtype == torch.float32, f"{logits.dtype}" assert len(logits.shape) == 3, f"{logits.shape}" assert logits.size(0) == 1, f"{logits.shape}" logits = logits.squeeze(0) @@ -508,32 +580,36 @@ def get_log_probs_and_entropy( with_entropy_grad = with_entropy and getattr(args, "entropy_coef", 0.0) != 0 # --- build full shifted-token target tensor --- - full_tokens = _build_shifted_tokens(T, device, unconcat_tokens, total_lengths, response_lengths, args.allgather_cp) - - # --- build top-p nucleus keep-mask (logprob only; entropy stays unmasked) --- - top_p_keep_mask = None - if top_p_token_ids is not None and top_p_token_offsets is not None: - top_p_keep_mask = _build_topp_keep_mask( - T, - logits.size(-1), - device, - top_p_token_ids, - top_p_token_offsets, - total_lengths, - response_lengths, - args.allgather_cp, - ) - - # --- compute on full [T,V] logits at once via calculate_log_probs_and_entropy --- - log_prob_full, entropy_full = calculate_log_probs_and_entropy( - logits, - full_tokens, - tp_group, - with_entropy=with_entropy, - with_entropy_grad=with_entropy_grad, - chunk_size=chunk_size, - log_prob_keep_mask=top_p_keep_mask, + full_tokens = _build_shifted_tokens( + T, + device, + unconcat_tokens, + total_lengths, + response_lengths, + qkv_format, + max_seq_lens, + args.allgather_cp, + padded_total_lengths, ) + + # --- compute log-probs; chunked-LM-head path fuses the matmul to avoid [T,V] --- + if lm_head_weight is not None: + log_prob_full, entropy_full = calculate_log_probs_and_entropy_fused( + logits, + lm_head_weight, + full_tokens, + tp_group, + with_entropy=with_entropy, + chunk_size=chunk_size, + ) + else: + log_prob_full, entropy_full = calculate_log_probs_and_entropy( + logits, + full_tokens, + tp_group, + with_entropy=with_entropy, + chunk_size=chunk_size, + ) log_prob_full = log_prob_full.squeeze(-1) # [T, 1] -> [T] # --- extract per-sample response portions --- @@ -542,7 +618,10 @@ def get_log_probs_and_entropy( entropy_full, total_lengths, response_lengths, + qkv_format, + max_seq_lens, args.allgather_cp, + padded_total_lengths, ) res = {"log_probs": log_probs_list} @@ -554,8 +633,11 @@ def get_log_probs_and_entropy( _allgather_cp_redistribute( res, logits_local_len=T, + args=args, total_lengths=total_lengths, response_lengths=response_lengths, + max_seq_lens=max_seq_lens, + padded_total_lengths=padded_total_lengths, ) return torch.empty((0,), device=device), res @@ -570,6 +652,8 @@ def get_values( response_lengths: list[int], with_entropy: bool = False, non_loss_data: bool = True, + max_seq_lens: list[int] | None = None, + padded_total_lengths: list[int] | None = None, ) -> dict[str, list[torch.Tensor]]: """Extract per-token value predictions over response tokens. @@ -597,6 +681,8 @@ def get_values( unconcat_tokens=unconcat_tokens, total_lengths=total_lengths, response_lengths=response_lengths, + max_seq_lens=max_seq_lens, + padded_total_lengths=padded_total_lengths, apply_temperature=False, ): assert logits_chunk.size(-1) == 1, f"{logits_chunk.shape}" @@ -610,8 +696,11 @@ def get_values( _allgather_cp_redistribute( res, logits_local_len=logits.size(1), + args=args, total_lengths=total_lengths, response_lengths=response_lengths, + max_seq_lens=max_seq_lens, + padded_total_lengths=padded_total_lengths, ) return torch.empty((0,), device=logits.device), res @@ -913,6 +1002,8 @@ def policy_loss_function( response_lengths = batch["response_lengths"] total_lengths = batch["total_lengths"] + max_seq_lens = batch.get("max_seq_lens", None) + padded_total_lengths = batch.get("padded_total_lengths", None) _, log_probs_and_entropy = get_log_probs_and_entropy( logits, @@ -921,7 +1012,8 @@ def policy_loss_function( total_lengths=total_lengths, response_lengths=response_lengths, with_entropy=True, - **get_rollout_top_p_logprob_kwargs(args, batch), + max_seq_lens=max_seq_lens, + padded_total_lengths=padded_total_lengths, ) log_probs = log_probs_and_entropy["log_probs"] @@ -937,16 +1029,18 @@ def policy_loss_function( full_log_probs = None full_old_log_probs = None if need_full_log_probs: + padded_iter = padded_total_lengths if padded_total_lengths is not None else [None] * len(total_lengths) full_log_probs = [ - all_gather_with_cp(log_prob, total_length, response_length) - for log_prob, total_length, response_length in zip( - log_probs, total_lengths, response_lengths, strict=False + all_gather_with_cp(log_prob, total_length, response_length, padded_total_length) + for log_prob, total_length, response_length, padded_total_length in zip( + log_probs, total_lengths, response_lengths, padded_iter, strict=False ) ] + padded_iter = padded_total_lengths if padded_total_lengths is not None else [None] * len(total_lengths) full_old_log_probs = [ - all_gather_with_cp(old_log_prob, total_length, response_length) - for old_log_prob, total_length, response_length in zip( - old_log_probs, total_lengths, response_lengths, strict=False + all_gather_with_cp(old_log_prob, total_length, response_length, padded_total_length) + for old_log_prob, total_length, response_length, padded_total_length in zip( + old_log_probs, total_lengths, response_lengths, padded_iter, strict=False ) ] @@ -1026,6 +1120,9 @@ def policy_loss_function( modified_response_masks, batch["rollout_mask_sums"], args.calculate_per_token_loss, + args.qkv_format, + max_seq_lens, + batch.get("padded_total_lengths", None), ) # Determine pg_loss reducer: use custom if specified, otherwise default @@ -1141,6 +1238,8 @@ def value_loss_function( unconcat_tokens=batch["unconcat_tokens"], total_lengths=batch["total_lengths"], response_lengths=batch["response_lengths"], + max_seq_lens=batch.get("max_seq_lens", None), + padded_total_lengths=batch.get("padded_total_lengths", None), ) values = torch.cat([value.flatten() for value in values["values"]], dim=0) @@ -1191,6 +1290,7 @@ def sft_loss_function( """ response_lengths = batch["response_lengths"] total_lengths = batch["total_lengths"] + lm_head_weight = _DEFERRED_LM_HEAD["weight"] if _chunked_lm_head_enabled() else None _, log_probs_and_entropy = get_log_probs_and_entropy( logits, @@ -1199,6 +1299,9 @@ def sft_loss_function( total_lengths=total_lengths, response_lengths=response_lengths, with_entropy=False, + max_seq_lens=batch.get("max_seq_lens", None), + padded_total_lengths=batch.get("padded_total_lengths", None), + lm_head_weight=lm_head_weight, ) log_probs = log_probs_and_entropy["log_probs"] @@ -1207,6 +1310,8 @@ def sft_loss_function( # make sure the gradient could backprop correctly. if log_probs.numel() == 0: + if lm_head_weight is not None: + loss += 0 * lm_head_weight.sum() loss += 0 * logits.sum() return ( @@ -1259,6 +1364,9 @@ def loss_function( batch["loss_masks"], batch["rollout_mask_sums"], args.calculate_per_token_loss, + args.qkv_format, + batch.get("max_seq_lens", None), + batch.get("padded_total_lengths", None), ) match args.loss_type: @@ -1318,3 +1426,56 @@ def loss_function( ), }, ) + + +def sft_precomputed_loss_function( + args: Namespace, + batch: RolloutBatch, + num_microbatches: int, + step_global_batch_size: int, + losses: torch.Tensor | tuple[torch.Tensor, torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor, dict[str, list[str] | torch.Tensor]]: + """Reduce per-token CE values returned by Megatron's output layer. + + VLM long-context SFT passes labels into the model so Megatron computes CE + before materializing full vocab logits. That avoids a large [T, vocab] + tensor while preserving Slime's per-token SFT scaling/reporting contract. + """ + if not args.calculate_per_token_loss: + raise ValueError("Precomputed VLM SFT loss requires --calculate-per-token-loss") + + if isinstance(losses, tuple): + losses, output_loss_mask = losses[:2] + loss_mask = output_loss_mask + else: + loss_mask = batch["full_loss_masks"] + + losses_view = losses.view(-1).float() + loss_mask_view = loss_mask.view(-1).float() + if losses_view.numel() != loss_mask_view.numel(): + # Qwen3-VL packed forward can pad hidden states for attention alignment, + # while linear CE returns only label-aligned losses. Treat tail padding + # as ignored if either side still carries it. + common_len = min(losses_view.numel(), loss_mask_view.numel()) + if loss_mask_view.numel() > common_len and loss_mask_view[common_len:].sum().ne(0): + raise RuntimeError( + "Precomputed SFT loss would drop nonzero tail loss_mask: " + f"losses={losses_view.numel()}, loss_mask={loss_mask_view.numel()}, common_len={common_len}" + ) + losses_view = losses_view[:common_len] + loss_mask_view = loss_mask_view[:common_len] + + local_loss = (losses_view * loss_mask_view).sum() + if args.allgather_cp and mpu.get_context_parallel_world_size() > 1: + local_loss = local_loss + 0 * losses.view(-1).float().sum() + num_tokens = sum(torch.clamp_min(loss_mask.sum(), 1) for loss_mask in batch["loss_masks"]) + loss = local_loss * mpu.get_context_parallel_world_size() + + return ( + loss, + num_tokens, + { + "keys": ["loss"], + "values": torch.stack([num_tokens.to(dtype=torch.float32), local_loss.detach().float()]), + }, + ) diff --git a/slime/backends/megatron_utils/megatron_bridge_compat.py b/slime/backends/megatron_utils/megatron_bridge_compat.py new file mode 100644 index 0000000000..8ce403642d --- /dev/null +++ b/slime/backends/megatron_utils/megatron_bridge_compat.py @@ -0,0 +1,46 @@ +"""Default-off compatibility shims for Megatron-Bridge runtime experiments.""" + +from __future__ import annotations + +import logging +import os + +logger = logging.getLogger(__name__) + + +def _env_flag(name: str) -> bool: + return os.environ.get(name, os.environ.get(name.lower(), "0")).lower() in {"1", "true", "yes", "on"} + + +def apply_megatron_bridge_compat_shims() -> None: + """Patch harmless import-time gaps between tested Bridge and image MCore. + + Newer Megatron-Bridge eagerly imports the Mamba provider from + ``megatron.bridge.models.__init__``. The QwenVL smoke does not use Mamba, + but the import can fail on older Megatron-Core pins before Qwen bridges are + registered. Keep the shim opt-in and make accidental Mamba use fail loudly. + """ + + if not _env_flag("SLIME_MBRIDGE_MAMBA_IMPORT_SHIM"): + return + + try: + import megatron.core.ssm.mamba_hybrid_layer_allocation as allocation + except Exception as exc: + logger.warning("Cannot import MCore Mamba allocation module for Bridge shim: %r", exc) + return + + if hasattr(allocation, "parse_hybrid_pattern"): + return + + def _parse_hybrid_pattern_unavailable(*args, **kwargs): + raise RuntimeError( + "SLIME_MBRIDGE_MAMBA_IMPORT_SHIM only bypasses Megatron-Bridge eager imports for non-Mamba models; " + "parse_hybrid_pattern is unavailable in this Megatron-Core image." + ) + + allocation.parse_hybrid_pattern = _parse_hybrid_pattern_unavailable + logger.warning( + "Installed SLIME_MBRIDGE_MAMBA_IMPORT_SHIM: added placeholder " + "megatron.core.ssm.mamba_hybrid_layer_allocation.parse_hybrid_pattern for non-Mamba Bridge imports." + ) diff --git a/slime/backends/megatron_utils/model.py b/slime/backends/megatron_utils/model.py index 1ad6cd7957..cae5e98ad3 100644 --- a/slime/backends/megatron_utils/model.py +++ b/slime/backends/megatron_utils/model.py @@ -33,14 +33,210 @@ from .checkpoint import load_checkpoint, save_checkpoint from .cp_utils import reduce_train_step_metrics -from .data import DataIterator, get_batch -from .loss import ROLLOUT_TOP_P_TOKEN_KEYS, get_rollout_top_p_logprob_kwargs, loss_function +from .data import ( + DataIterator, + get_batch, + has_multimodal_train_inputs, + qwen_vl_text_fastpath_requires_unsplit_input, + qwen_vl_unsplit_only_with_mm, +) +from .initialize import get_use_gloo_process_groups +from .loss import _build_shifted_tokens, loss_function, set_deferred_lm_head_weight, sft_precomputed_loss_function from .model_provider import get_model_provider_func from .stateless_adam import StatelessAdam logger = logging.getLogger(__name__) +_DEFER_ACTIVE = {"on": False} +_QWENVL_FORWARD_DEBUG_COUNT = 0 + + +def _qwen_vl_external_sft_loss() -> bool: + return os.environ.get("SLIME_QWENVL_EXTERNAL_SFT_LOSS", "0").lower() in {"1", "true", "yes", "on"} + + +def _qwen_vl_text_language_fastpath() -> bool: + return os.environ.get("SLIME_QWENVL_TEXT_LANGUAGE_FASTPATH", "0").lower() in {"1", "true", "yes", "on"} + + +def _qwen_vl_text_language_fastpath_safe() -> bool: + return _qwen_vl_text_language_fastpath() and not qwen_vl_text_fastpath_requires_unsplit_input() + + +def _qwen_vl_text_fastpath_local_mrope() -> bool: + return os.environ.get("SLIME_QWENVL_TEXT_FASTPATH_LOCAL_MROPE", "1").lower() in {"1", "true", "yes", "on"} + + +def _env_flag(name: str) -> bool: + return os.environ.get(name, os.environ.get(name.lower(), "0")).lower() in {"1", "true", "yes", "on"} + + +def _env_int(name: str, default: int) -> int: + try: + return int(os.environ.get(name, os.environ.get(name.lower(), str(default)))) + except ValueError: + return default + + +def _distributed_rank() -> int: + try: + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + return dist.get_rank() + except Exception: + pass + return int(os.environ.get("RANK", "0")) + + +def _debug_tensor_summary(value): + if torch.is_tensor(value): + return { + "shape": tuple(value.shape), + "dtype": str(value.dtype).replace("torch.", ""), + "device": str(value.device), + "numel": value.numel(), + } + return None + + +def _debug_packed_seq_params(value): + if type(value).__name__ != "PackedSeqParams": + return None + summary = { + "qkv_format": getattr(value, "qkv_format", None), + "max_seqlen_q": getattr(value, "max_seqlen_q", None), + "max_seqlen_kv": getattr(value, "max_seqlen_kv", None), + } + for name in ("cu_seqlens_q", "cu_seqlens_kv", "cu_seqlens_q_padded", "cu_seqlens_kv_padded"): + tensor = getattr(value, name, None) + if not torch.is_tensor(tensor): + summary[name] = None + continue + vals = tensor.detach().cpu() + diffs = vals[1:] - vals[:-1] if vals.numel() > 1 else vals.new_empty((0,)) + summary[name] = { + "shape": tuple(tensor.shape), + "dtype": str(tensor.dtype).replace("torch.", ""), + "device": str(tensor.device), + "head": vals[:8].tolist(), + "tail": vals[-8:].tolist(), + "diff_min": int(diffs.min().item()) if diffs.numel() else None, + "diff_max": int(diffs.max().item()) if diffs.numel() else None, + } + return summary + + +def _debug_value_summary(value, depth: int = 0): + if depth > 2: + return type(value).__name__ + tensor_summary = _debug_tensor_summary(value) + if tensor_summary is not None: + return tensor_summary + packed_summary = _debug_packed_seq_params(value) + if packed_summary is not None: + return {"type": "PackedSeqParams", **packed_summary} + if isinstance(value, dict): + return {str(key): _debug_value_summary(val, depth + 1) for key, val in list(value.items())[:12]} + if isinstance(value, (list, tuple)): + return [_debug_value_summary(item, depth + 1) for item in value[:4]] + if value is None or isinstance(value, (bool, int, float, str)): + return value + return type(value).__name__ + + +def _maybe_log_qwenvl_forward_debug( + args: Namespace, + batch: dict, + forward_kwargs: dict, + *, + has_mm_inputs: bool, + needs_unsplit: bool, + use_unsplit: bool, + use_precomputed_sft_loss: bool = False, +) -> None: + global _QWENVL_FORWARD_DEBUG_COUNT + + if not _env_flag("SLIME_QWENVL_FORWARD_SHAPE_DEBUG"): + return + + rank_filter = _env_int("SLIME_QWENVL_FORWARD_SHAPE_DEBUG_RANK", 0) + rank = _distributed_rank() + if rank_filter >= 0 and rank != rank_filter: + return + + limit = _env_int("SLIME_QWENVL_FORWARD_SHAPE_DEBUG_LIMIT", 16) + if _QWENVL_FORWARD_DEBUG_COUNT >= limit: + return + _QWENVL_FORWARD_DEBUG_COUNT += 1 + + try: + logger.info( + "QwenVL forward shape debug: sample=%s rank=%s has_mm_inputs=%s " + "uses_unsplit_forward=%s qwen_vl_unsplit_only_with_mm=%s needs_unsplit=%s " + "use_unsplit=%s use_precomputed_sft_loss=%s qkv_format=%s allgather_cp=%s " + "input_ids=%s position_ids=%s attention_mask=%s loss_mask=%s labels=%s " + "packed_seq_params=%s batch_packed_seq_params=%s vlm_packed_seq_params=%s " + "tokens=%s unsplit_tokens=%s full_loss_masks=%s multimodal_train_inputs=%s " + "max_seq_lens=%s padded_total_lengths=%s", + _QWENVL_FORWARD_DEBUG_COUNT, + rank, + has_mm_inputs, + getattr(args, "uses_unsplit_forward", False), + qwen_vl_unsplit_only_with_mm(), + needs_unsplit, + use_unsplit, + use_precomputed_sft_loss, + getattr(args, "qkv_format", None), + getattr(args, "allgather_cp", None), + _debug_value_summary(forward_kwargs.get("input_ids")), + _debug_value_summary(forward_kwargs.get("position_ids")), + _debug_value_summary(forward_kwargs.get("attention_mask")), + _debug_value_summary(forward_kwargs.get("loss_mask")), + _debug_value_summary(forward_kwargs.get("labels")), + _debug_value_summary(forward_kwargs.get("packed_seq_params")), + _debug_value_summary(batch.get("packed_seq_params")), + _debug_value_summary(batch.get("vlm_packed_seq_params")), + _debug_value_summary(batch.get("tokens")), + _debug_value_summary(batch.get("unsplit_tokens")), + _debug_value_summary(batch.get("full_loss_masks")), + _debug_value_summary(batch.get("multimodal_train_inputs")), + _debug_value_summary(batch.get("max_seq_lens")), + _debug_value_summary(batch.get("padded_total_lengths")), + ) + except Exception: + logger.exception("Failed to log QwenVL forward shape debug") + + +def _unwrap_to_core_model(model): + while hasattr(model, "module"): + model = model.module + return model + + +def _ensure_deferred_output_layer(model) -> bool: + import types + + core = _unwrap_to_core_model(model) + output_layer = getattr(core, "output_layer", None) + if output_layer is None or getattr(output_layer, "weight", None) is None: + return False + + set_deferred_lm_head_weight(output_layer.weight) + if not getattr(output_layer, "_slime_deferred", False): + output_layer._slime_orig_forward = output_layer.forward + + def _deferred_forward(self, input_, *args, **kwargs): + if _DEFER_ACTIVE["on"]: + return input_, None + return self._slime_orig_forward(input_, *args, **kwargs) + + output_layer.forward = types.MethodType(_deferred_forward, output_layer) + output_layer._slime_deferred = True + return True + + def _disable_tqdm_for_non_main_rank() -> bool: return not ( mpu.get_data_parallel_rank(with_context_parallel=True) == 0 @@ -301,19 +497,11 @@ def setup_model_and_optimizer( config = OptimizerConfig(**kwargs) config.timers = None - if args.use_stateless_adam: - assert config.optimizer == "adam", "Stateless Adam only supports --optimizer adam." - assert args.no_save_optim, "Stateless Adam does not save Adam moment states. Please set --no-save-optim." - - optimizer_context = _patch_megatron_adam(StatelessAdam) if args.use_stateless_adam else nullcontext() - with optimizer_context: - optimizer = get_megatron_optimizer( - config=config, - model_chunks=model, - use_gloo_process_groups=args.enable_gloo_process_groups, - ) - if args.use_stateless_adam: - _disable_distributed_optimizer_state_initialization(optimizer) + optimizer = get_megatron_optimizer( + config=config, + model_chunks=model, + use_gloo_process_groups=get_use_gloo_process_groups(args), + ) opt_param_scheduler = get_optimizer_param_scheduler(args, optimizer) return model, optimizer, opt_param_scheduler @@ -420,29 +608,68 @@ def forward_step( packed_seq_params = batch["packed_seq_params"] total_lengths = batch["total_lengths"] response_lengths = batch["response_lengths"] + + has_mm_inputs = has_multimodal_train_inputs(batch.get("multimodal_train_inputs", None)) + uses_unsplit_forward = getattr(args, "uses_unsplit_forward", False) + needs_unsplit = has_mm_inputs or ( + uses_unsplit_forward + and (not qwen_vl_unsplit_only_with_mm() or qwen_vl_text_fastpath_requires_unsplit_input()) + ) + text_language_fastpath = ( + uses_unsplit_forward + and not needs_unsplit + and _qwen_vl_text_language_fastpath_safe() + and batch.get("packed_seq_params") is not None + ) + mm_kwargs = batch["multimodal_train_inputs"] if has_mm_inputs else {} + use_unsplit = needs_unsplit and "unsplit_tokens" in batch + + position_ids = None + if getattr(args, "position_embedding_type", "rope") == "mrope" and not use_unsplit: + from .mrope_utils import build_mrope_position_ids + + position_ids = build_mrope_position_ids( + batch, + local_thd_cp=text_language_fastpath and _qwen_vl_text_fastpath_local_mrope(), + ) + forward_kwargs = { - "input_ids": tokens, - "position_ids": None, + "input_ids": batch["unsplit_tokens"] if use_unsplit else tokens, + "position_ids": position_ids, "attention_mask": None, "labels": None, - "packed_seq_params": packed_seq_params, + "packed_seq_params": None if use_unsplit else packed_seq_params, "loss_mask": batch["full_loss_masks"], } - if batch["multimodal_train_inputs"] is not None: - forward_kwargs.update(batch["multimodal_train_inputs"]) - output_tensor = model(**forward_kwargs) - output_kwargs = { - "args": args, - "unconcat_tokens": unconcat_tokens, - "total_lengths": total_lengths, - "response_lengths": response_lengths, - "with_entropy": args.use_rollout_entropy, - } - if use_rollout_top_p_replay: - output_kwargs.update(get_rollout_top_p_logprob_kwargs(args, batch)) + if needs_unsplit and "vlm_packed_seq_params" in batch: + forward_kwargs["attention_mask"] = batch["unsplit_attention_mask"] + forward_kwargs["packed_seq_params"] = batch["vlm_packed_seq_params"] + forward_kwargs["loss_mask"] = None + + if has_mm_inputs: + forward_kwargs.update(mm_kwargs) - return output_tensor, partial(f, **output_kwargs) + _maybe_log_qwenvl_forward_debug( + args, + batch, + forward_kwargs, + has_mm_inputs=has_mm_inputs, + needs_unsplit=needs_unsplit, + use_unsplit=use_unsplit, + ) + output_tensor = model(**forward_kwargs) + + 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, + max_seq_lens=batch.get("max_seq_lens", None), + padded_total_lengths=batch.get("padded_total_lengths", None), + ) # Turn on evaluation mode which disables dropout. for model_module in model: @@ -603,9 +830,16 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p old_stage = os.environ["ROUTING_REPLAY_STAGE"] os.environ["ROUTING_REPLAY_STAGE"] = "replay_forward" + use_precomputed_sft_loss = False + if return_schedule_plan: assert not args.enable_mtp_training, "MTP training should not be enabled when using combined 1f1b" position_ids = None + if getattr(args, "position_embedding_type", "rope") == "mrope": + from .mrope_utils import build_mrope_position_ids + + position_ids = build_mrope_position_ids(batch) + output_tensor = model.build_schedule_plan( input_ids=batch["tokens"], position_ids=position_ids, @@ -615,27 +849,105 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p loss_mask=batch["full_loss_masks"], ) else: + has_mm_inputs = has_multimodal_train_inputs(batch.get("multimodal_train_inputs", None)) + uses_unsplit_forward = getattr(args, "uses_unsplit_forward", False) + needs_unsplit = has_mm_inputs or ( + uses_unsplit_forward + and (not qwen_vl_unsplit_only_with_mm() or qwen_vl_text_fastpath_requires_unsplit_input()) + ) + text_language_fastpath = ( + uses_unsplit_forward + and not needs_unsplit + and _qwen_vl_text_language_fastpath_safe() + and batch.get("packed_seq_params") is not None + ) + use_unsplit = needs_unsplit and "unsplit_tokens" in batch + use_precomputed_sft_loss = ( + args.loss_type == "sft_loss" + and (needs_unsplit or text_language_fastpath) + and getattr(args, "calculate_per_token_loss", False) + and not _qwen_vl_external_sft_loss() + ) + + position_ids = None + if getattr(args, "position_embedding_type", "rope") == "mrope" and not use_unsplit: + from .mrope_utils import build_mrope_position_ids + + position_ids = build_mrope_position_ids( + batch, + local_thd_cp=text_language_fastpath and _qwen_vl_text_fastpath_local_mrope(), + ) + forward_kwargs = { - "input_ids": batch["tokens"], - "position_ids": None, + "input_ids": batch["unsplit_tokens"] if use_unsplit else batch["tokens"], + "position_ids": position_ids, "attention_mask": None, "labels": None, - "packed_seq_params": batch["packed_seq_params"], + "packed_seq_params": None if use_unsplit else batch["packed_seq_params"], "loss_mask": batch["full_loss_masks"], } - if batch["multimodal_train_inputs"] is not None: - forward_kwargs.update(batch["multimodal_train_inputs"]) + if needs_unsplit and "vlm_packed_seq_params" in batch: + forward_kwargs["attention_mask"] = batch["unsplit_attention_mask"] + forward_kwargs["packed_seq_params"] = batch["vlm_packed_seq_params"] + if not use_precomputed_sft_loss: + forward_kwargs["loss_mask"] = None + + if use_precomputed_sft_loss: + shifted_labels = _build_shifted_tokens( + batch["tokens"].numel(), + batch["tokens"].device, + batch["unconcat_tokens"], + batch["total_lengths"], + batch["response_lengths"], + args.qkv_format, + batch.get("max_seq_lens", None), + args.allgather_cp, + batch.get("padded_total_lengths", None), + ).view_as(batch["tokens"]) + forward_kwargs["labels"] = shifted_labels.masked_fill( + batch["full_loss_masks"].view_as(shifted_labels).le(0), + -100, + ) + forward_kwargs["loss_mask"] = batch["full_loss_masks"] if args.enable_mtp_training: forward_kwargs["mtp_kwargs"] = {"mtp_labels": batch["tokens"]} - output_tensor = model(**forward_kwargs) + if has_mm_inputs: + forward_kwargs.update(batch["multimodal_train_inputs"]) + + _maybe_log_qwenvl_forward_debug( + args, + batch, + forward_kwargs, + has_mm_inputs=has_mm_inputs, + needs_unsplit=needs_unsplit, + use_unsplit=use_unsplit, + use_precomputed_sft_loss=use_precomputed_sft_loss, + ) + if ( + os.environ.get("CHUNKED_LM_HEAD", "0") == "1" + and args.loss_type == "sft_loss" + and not use_precomputed_sft_loss + and _ensure_deferred_output_layer(model) + ): + _DEFER_ACTIVE["on"] = True + try: + output_tensor = model(**forward_kwargs) + finally: + _DEFER_ACTIVE["on"] = False + else: + output_tensor = model(**forward_kwargs) 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, step_global_batch_size) + if not return_schedule_plan and use_precomputed_sft_loss: + loss_func = sft_precomputed_loss_function + else: + loss_func = loss_function + return output_tensor, partial(loss_func, args, batch, num_microbatches, step_global_batch_size) # Forward pass. forward_backward_func = get_forward_backward_func() diff --git a/slime/backends/megatron_utils/model_provider.py b/slime/backends/megatron_utils/model_provider.py index 090b692f18..950ca4b5bd 100644 --- a/slime/backends/megatron_utils/model_provider.py +++ b/slime/backends/megatron_utils/model_provider.py @@ -1,6 +1,7 @@ # Adapt from https://github.com/NVIDIA/Megatron-LM/blob/b1efb3c7126ef7615e8c333432d76e08038e17ff/pretrain_gpt.py import argparse import inspect +import logging import re from contextlib import nullcontext from typing import Literal @@ -20,6 +21,35 @@ from slime.utils.megatron_bridge_utils import patch_auto_bridge_hf_config from slime.utils.misc import load_function +logger = logging.getLogger(__name__) + + +def _maybe_apply_megatron_bridge_compat_shims() -> None: + try: + from .megatron_bridge_compat import apply_megatron_bridge_compat_shims + except Exception as exc: + logger.warning("Failed to import Megatron-Bridge compat shims: %r", exc) + return + apply_megatron_bridge_compat_shims() + + +def _maybe_mark_unsplit_forward(args: argparse.Namespace, model: torch.nn.Module) -> None: + """Mark bridge Qwen-VL models that expect unsplit input under CP. + + Qwen3VLModel performs the CP/SP split inside its own forward after vision + embedding and packed-sequence preprocessing. This also applies to + text-only Qwen3.5/Qwen3.6 samples sharing the same bridge model class. + """ + try: + from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model import Qwen3VLModel + except ImportError: + return + + if isinstance(model, Qwen3VLModel): + if not getattr(args, "uses_unsplit_forward", False): + logger.info("Detected Qwen3VLModel; enabling args.uses_unsplit_forward.") + args.uses_unsplit_forward = True + # Adapt from https://github.com/volcengine/verl/blob/c3b20575d2bc815fcccd84bddb4c0401fc4b632b/verl/models/llama/megatron/layers/parallel_linear.py#L82 class LinearForLastLayer(torch.nn.Linear): @@ -80,47 +110,117 @@ def wrapped_model_provider( model.output_layer = LinearForLastLayer( input_size=model.config.hidden_size, output_size=1, config=model.config ) + _maybe_mark_unsplit_forward(args, model) return model return wrapped_model_provider if args.megatron_to_hf_mode == "bridge": + _maybe_apply_megatron_bridge_compat_shims() from megatron.bridge import AutoBridge import slime_plugins.megatron_bridge # noqa: F401 # register custom bridges bridge = patch_auto_bridge_hf_config(AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)) provider = bridge.to_megatron_provider(load_weights=False) - # TODO: we should not manually set this... - provider.tensor_model_parallel_size = args.tensor_model_parallel_size - provider.pipeline_model_parallel_size = args.pipeline_model_parallel_size - provider.expert_model_parallel_size = args.expert_model_parallel_size - provider.expert_tensor_parallel_size = args.expert_tensor_parallel_size - provider.sequence_parallel = args.sequence_parallel - provider.context_parallel_size = args.context_parallel_size - provider.variable_seq_lengths = args.variable_seq_lengths - if hasattr(args, "moe_token_dispatcher_type"): - provider.moe_token_dispatcher_type = args.moe_token_dispatcher_type + + # Keep Megatron Bridge providers in sync with Slime/Megatron args for + # long-context VLM SFT. Bridge defaults are often conservative and may + # disagree with the already-validated launch topology. + bridge_keys = [ + "attention_backend", + "tensor_model_parallel_size", + "pipeline_model_parallel_size", + "context_parallel_size", + "expert_model_parallel_size", + "expert_tensor_parallel_size", + "sequence_parallel", + "variable_seq_lengths", + "attention_softmax_in_fp32", + "bias_dropout_fusion", + "apply_rope_fusion", + "recompute_granularity", + "recompute_method", + "recompute_num_layers", + "distribute_saved_activations", + "moe_router_load_balancing_type", + "moe_router_dtype", + "moe_aux_loss_coeff", + "moe_token_dispatcher_type", + "moe_shared_expert_overlap", + "moe_enable_deepep", + "moe_flex_dispatcher_backend", + "freeze_language_model", + "freeze_vision_model", + "freeze_vision_projection", + "vision_dp_when_cp", + "vision_dp_when_tp", + "calculate_per_token_loss", + "num_layers", + "moe_layer_freq", + ] + + args_dict = vars(args) + for attr in bridge_keys: + if attr not in args_dict or not hasattr(provider, attr): + continue + old_val = getattr(provider, attr) + new_val = args_dict[attr] + if old_val != new_val: + logger.info(f"Override provider.{attr}: {old_val!r} -> {new_val!r}") + setattr(provider, attr, new_val) + + if getattr(args, "loss_type", None) == "sft_loss" and hasattr(provider, "cross_entropy_loss_fusion"): + # Long-context VLM SFT passes labels into Megatron so the model can + # return per-token CE directly. The native CE path still + # materializes [T, vocab] logits; linear CE avoids that tensor. + if not provider.cross_entropy_loss_fusion: + logger.info("Override provider.cross_entropy_loss_fusion: False -> True") + provider.cross_entropy_loss_fusion = True + if hasattr(provider, "cross_entropy_fusion_impl"): + old_val = getattr(provider, "cross_entropy_fusion_impl") + if old_val != "linear": + logger.info(f"Override provider.cross_entropy_fusion_impl: {old_val!r} -> 'linear'") + provider.cross_entropy_fusion_impl = "linear" + + if hasattr(provider, "mtp_num_layers"): + # Qwen3.5/3.6 HF configs may advertise MTP layers for speculative + # decoding. Plain SFT does not train MTP, and loading those extra + # Megatron params from a base HF checkpoint can produce missing + # bridge tasks. Only keep them when MTP training is explicitly on. + new_val = getattr(args, "mtp_num_layers", None) if getattr(args, "enable_mtp_training", False) else None + old_val = getattr(provider, "mtp_num_layers") + if old_val != new_val: + logger.info(f"Override provider.mtp_num_layers: {old_val!r} -> {new_val!r}") + provider.mtp_num_layers = new_val + if getattr(args, "decoder_first_pipeline_num_layers", None) is not None: provider.num_layers_in_first_pipeline_stage = args.decoder_first_pipeline_num_layers if getattr(args, "decoder_last_pipeline_num_layers", None) is not None: provider.num_layers_in_last_pipeline_stage = args.decoder_last_pipeline_num_layers + if getattr(args, "fp16", False): + provider.fp16 = True + provider.bf16 = False + provider.params_dtype = torch.float16 + elif getattr(args, "bf16", False): + provider.fp16 = False + provider.bf16 = True + provider.params_dtype = torch.bfloat16 provider.finalize() - if role == "critic": - _original_provide = provider.provide + _original_provide = provider.provide - def _critic_provide(pre_process=True, post_process=True, vp_stage=None): - model = _original_provide(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage) + def _provide(pre_process=True, post_process=True, vp_stage=None): + model = _original_provide(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage) + _maybe_mark_unsplit_forward(args, model) + if role == "critic": if post_process: model.output_layer = LinearForLastLayer( input_size=model.config.hidden_size, output_size=1, config=model.config ) - return model - - return _critic_provide + return model - return provider.provide + return _provide def model_provider(pre_process: bool = True, post_process: bool = True, vp_stage: int | None = None) -> GPTModel: """Builds the model. @@ -153,6 +253,7 @@ def model_provider(pre_process: bool = True, post_process: bool = True, vp_stage model.output_layer = LinearForLastLayer( input_size=config.hidden_size, output_size=1, config=config ) + _maybe_mark_unsplit_forward(args, model) return model transformer_layer_spec = result else: @@ -237,6 +338,7 @@ def model_provider(pre_process: bool = True, post_process: bool = True, vp_stage if post_process and role == "critic": model.output_layer = LinearForLastLayer(input_size=config.hidden_size, output_size=1, config=config) + _maybe_mark_unsplit_forward(args, model) return model return model_provider @@ -269,13 +371,58 @@ def get_model_provider_func(args, role="actor"): return wrap_model_provider_with_freeze(_get_model_provider_func(args, role), args) +def _log_freeze_summary(model: GPTModel, args: argparse.Namespace, matched_patterns: dict[str, list[str]]) -> None: + try: + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized() and dist.get_rank() != 0: + return + except Exception: + pass + + if not ( + getattr(args, "only_train_params_name_list", None) + or getattr(args, "freeze_params_name_list", None) + ): + return + + total_tensors = 0 + total_numel = 0 + trainable_tensors = 0 + trainable_numel = 0 + for _, param in model.named_parameters(): + total_tensors += 1 + total_numel += param.numel() + if param.requires_grad: + trainable_tensors += 1 + trainable_numel += param.numel() + + match_summary = { + pattern: {"count": len(names), "examples": names[:8]} + for pattern, names in matched_patterns.items() + } + logger.info( + "Freeze parameter summary: only_train=%s freeze=%s total_tensors=%s total_numel=%s " + "trainable_tensors=%s trainable_numel=%s matched=%s", + getattr(args, "only_train_params_name_list", None), + getattr(args, "freeze_params_name_list", None), + total_tensors, + total_numel, + trainable_tensors, + trainable_numel, + match_summary, + ) + + def freeze_model_params(model: GPTModel, args: argparse.Namespace): + matched_patterns: dict[str, list[str]] = {} if getattr(args, "only_train_params_name_list", None): for name, param in model.named_parameters(): param.requires_grad = False for pattern in args.only_train_params_name_list: if re.search(pattern, name): param.requires_grad = True + matched_patterns.setdefault(f"only_train:{pattern}", []).append(name) break if getattr(args, "freeze_params_name_list", None): @@ -283,4 +430,7 @@ def freeze_model_params(model: GPTModel, args: argparse.Namespace): for pattern in args.freeze_params_name_list: if re.search(pattern, name): param.requires_grad = False + matched_patterns.setdefault(f"freeze:{pattern}", []).append(name) break + + _log_freeze_summary(model, args, matched_patterns) diff --git a/slime/backends/megatron_utils/mrope_utils.py b/slime/backends/megatron_utils/mrope_utils.py new file mode 100644 index 0000000000..f4dea8baf7 --- /dev/null +++ b/slime/backends/megatron_utils/mrope_utils.py @@ -0,0 +1,63 @@ +"""mRoPE helpers for text-only Megatron training batches.""" + +import torch + + +def _build_local_thd_cp_position_ids(batch): + from megatron.core import mpu + + cp_size = mpu.get_context_parallel_world_size() + cp_rank = mpu.get_context_parallel_rank() + if cp_size == 1: + return None + + device = batch["tokens"].device + parts = [] + for tokens in batch["unconcat_tokens"]: + token_len = int(tokens.size(0)) + chunk_size = (token_len + 2 * cp_size - 1) // (2 * cp_size) + for start in ( + cp_rank * chunk_size, + (2 * cp_size - cp_rank - 1) * chunk_size, + ): + position = torch.arange(start, start + chunk_size, device=device, dtype=torch.long) + parts.append(torch.where(position < token_len, position, torch.zeros_like(position))) + + if not parts: + return None + + position = torch.cat(parts) + target_len = int(batch["tokens"].numel()) + if position.numel() > target_len: + raise ValueError(f"local mRoPE position length {position.numel()} exceeds token length {target_len}") + if position.numel() < target_len: + position = torch.cat( + [ + position, + torch.zeros(target_len - position.numel(), device=device, dtype=position.dtype), + ] + ) + return position.view(1, 1, target_len).expand(3, 1, target_len).contiguous() + + +def build_mrope_position_ids(batch, *, local_thd_cp: bool = False): + """Return position ids of shape [3, 1, max_seqlen] for text-only mRoPE. + + For THD, mirror the packed sequence max length so Megatron can slice per + document. For BSHD, packed sequence params are intentionally absent, so + use the padded sequence length from the token tensor. Text-only data uses + identical temporal/height/width axes. + """ + if local_thd_cp: + position_ids = _build_local_thd_cp_position_ids(batch) + if position_ids is not None: + return position_ids + + packed_seq_params = batch["packed_seq_params"] + if packed_seq_params is None: + max_seqlen = int(batch["tokens"].shape[-1]) + else: + max_seqlen = int(packed_seq_params.max_seqlen_q) + + position = torch.arange(max_seqlen, device=batch["tokens"].device, dtype=torch.long) + return position.view(1, 1, max_seqlen).expand(3, 1, max_seqlen).contiguous() diff --git a/slime/backends/megatron_utils/qwen_vl_packed_gdn.py b/slime/backends/megatron_utils/qwen_vl_packed_gdn.py new file mode 100644 index 0000000000..4879447e38 --- /dev/null +++ b/slime/backends/megatron_utils/qwen_vl_packed_gdn.py @@ -0,0 +1,527 @@ +"""Packed-sequence GDN patch for Qwen3.5/3.6 VL Bridge models. + +Megatron-Bridge's Qwen3.5 VL providers build linear-attention layers with +Megatron-Core ``GatedDeltaNet``. Slime's Qwen3.5 text plugin already supports +packed varlen GDN via FLA/FlashQLA, but Bridge VLM models bypass that plugin and +hit Megatron-Core's packed-sequence guard. This patch keeps the Bridge weight +layout intact and only fills in the packed forward path. +""" + +from __future__ import annotations + +import logging +import os +import time + +import torch +import torch.nn.functional as F + +logger = logging.getLogger(__name__) + + +def _env_flag(name: str) -> bool: + return os.environ.get(name, "0").lower() in {"1", "true", "yes", "on"} + + +def _env_int(name: str, default: int) -> int: + try: + return int(os.environ.get(name, str(default))) + except ValueError: + return default + + +def _get_rank() -> int: + try: + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + return dist.get_rank() + except Exception: + pass + return int(os.environ.get("RANK", "0")) + + +def _gdn_timing_enabled(module) -> bool: + if not _env_flag("SLIME_QWENVL_GDN_TIMING"): + return False + rank_filter = _env_int("SLIME_QWENVL_GDN_TIMING_RANK", 0) + return rank_filter < 0 or _get_rank() == rank_filter + + +def _gdn_timing_should_sample(module) -> bool: + if not _gdn_timing_enabled(module): + return False + count = getattr(module, "_slime_gdn_timing_count", 0) + 1 + module._slime_gdn_timing_count = count + limit = _env_int("SLIME_QWENVL_GDN_TIMING_LIMIT", 64) + interval = max(1, _env_int("SLIME_QWENVL_GDN_TIMING_INTERVAL", 32)) + return count <= limit and (count == 1 or count % interval == 0) + + +class _CudaStepTimer: + def __init__(self, enabled: bool): + self.enabled = enabled and torch.cuda.is_available() + self.events: list[tuple[str, torch.cuda.Event, torch.cuda.Event]] = [] + self.start_wall = time.perf_counter() + + def run(self, name, fn): + if not self.enabled: + return fn() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + result = fn() + end.record() + self.events.append((name, start, end)) + return result + + def summary(self) -> dict[str, float]: + if not self.enabled: + return {"wall_ms": (time.perf_counter() - self.start_wall) * 1000.0} + torch.cuda.synchronize() + summary = {name: start.elapsed_time(end) for name, start, end in self.events} + summary["total_ms"] = sum(summary.values()) + return summary + + +def _get_packed_cu_seqlens(packed_seq_params): + return ( + packed_seq_params.cu_seqlens_q_padded + if packed_seq_params.cu_seqlens_q_padded is not None + else packed_seq_params.cu_seqlens_q + ) + + +def _local_cu_seqlens_for_kernel(module, cu_seqlens: torch.Tensor, total_seq_len: int) -> torch.Tensor: + """Convert global packed cu_seqlens to the local physical sequence layout. + + Bridge Qwen3-VL builds THD packed params with padded *global* sequence + boundaries. After CP preprocessing, each rank holds ``padded_len / cp`` + tokens per sample. FLA kernels operate on this local tensor and will + illegal-access if they receive the full global boundary. + """ + expected_total = int(cu_seqlens[-1].detach().item()) if cu_seqlens.numel() > 0 else total_seq_len + if expected_total == total_seq_len: + return cu_seqlens + + cp_size = max(int(getattr(module, "cp_size", 1)), 1) + if expected_total % cp_size == 0 and expected_total // cp_size == total_seq_len: + return torch.div(cu_seqlens, cp_size, rounding_mode="floor").to(dtype=cu_seqlens.dtype) + + if not getattr(module, "_slime_gdn_cu_length_warned", False): + logger.warning( + "Packed GDN cu_seqlens mismatch: cu_seqlens[-1]=%s physical_seq_len=%s cp_size=%s; " + "falling back to one local range.", + expected_total, + total_seq_len, + cp_size, + ) + module._slime_gdn_cu_length_warned = True + return torch.tensor([0, total_seq_len], dtype=cu_seqlens.dtype, device=cu_seqlens.device) + + +def _get_gated_delta_rule(module): + backend = "fla" + try: + from megatron.training.global_vars import get_args + + backend = getattr(get_args(), "qwen_gdn_backend", backend) + except Exception: + pass + + if getattr(module, "_slime_gdn_backend", None) == backend and hasattr(module, "_slime_gated_delta_rule"): + return module._slime_gated_delta_rule + + try: + from slime_plugins.models.qwen_gdn_backend import get_chunk_gated_delta_rule + + rule = get_chunk_gated_delta_rule(backend) + except Exception as exc: + if backend != "fla": + logger.warning("Qwen GDN backend %s is unavailable: %s; falling back to fla.", backend, exc) + backend = "fla" + rule = module.gated_delta_rule + + module._slime_gdn_backend = backend + module._slime_gated_delta_rule = rule + if not getattr(module, "_slime_gdn_backend_logged", False): + logger.info("Qwen VL packed GDN backend selected: %s", backend) + module._slime_gdn_backend_logged = True + return rule + + +def _packed_ranges_from_cu_seqlens(cu_seqlens: torch.Tensor, total_seq_len: int) -> list[tuple[int, int]]: + """Return clipped packed sequence ranges covering the physical qkv tensor.""" + values = cu_seqlens.detach().to(device="cpu", dtype=torch.long).tolist() + if not values: + return [(0, total_seq_len)] + if values[0] != 0: + values = [0, *values] + + ranges: list[tuple[int, int]] = [] + prev = 0 + for value in values[1:]: + end = min(max(int(value), prev), total_seq_len) + if end > prev: + ranges.append((prev, end)) + prev = end + + if prev < total_seq_len: + ranges.append((prev, total_seq_len)) + return ranges or [(0, total_seq_len)] + + +def _range_stats(ranges: list[tuple[int, int]]) -> dict[str, int | float]: + lengths = [end - start for start, end in ranges] + if not lengths: + return {"ranges": 0, "min": 0, "max": 0, "mean": 0.0} + return { + "ranges": len(lengths), + "min": min(lengths), + "max": max(lengths), + "mean": sum(lengths) / len(lengths), + } + + +def _safe_packed_causal_conv1d(module, qkv, conv1d_weight, conv1d_bias, cu_seqlens): + """Torch fallback for packed Qwen GDN conv. + + FLA's Triton causal_conv1d varlen path can illegal-access on short padded + packed sequences. The grouped torch conv path matches Megatron-Core's + deterministic implementation and resets convolution state at each packed + sequence boundary. + """ + total_seq_len = qkv.shape[1] + outputs = [] + groups = module.conv_dim_local_tp // module.cp_size + for start, end in _packed_ranges_from_cu_seqlens(cu_seqlens, total_seq_len): + segment = qkv[:, start:end, :].transpose(1, 2).contiguous() + conv_out = F.conv1d( + input=segment, + weight=conv1d_weight, + bias=conv1d_bias, + stride=module.conv1d.stride, + padding=module.conv1d.padding, + dilation=module.conv1d.dilation, + groups=groups, + ) + conv_out = module.act_fn(conv_out[..., : end - start]) + outputs.append(conv_out.transpose(1, 2)) + + return torch.cat(outputs, dim=1).contiguous() + + +def _as_1d_int(value) -> int: + if isinstance(value, tuple): + return int(value[0]) + return int(value) + + +def _batched_safe_packed_causal_conv1d(module, qkv, conv1d_weight, conv1d_bias, cu_seqlens): + """Run the safe packed conv as one torch conv with zero gaps between ranges.""" + total_seq_len = qkv.shape[1] + ranges = _packed_ranges_from_cu_seqlens(cu_seqlens, total_seq_len) + if len(ranges) <= 1: + return _safe_packed_causal_conv1d(module, qkv, conv1d_weight, conv1d_bias, cu_seqlens) + + kernel_size = _as_1d_int(module.conv1d.kernel_size) + dilation = _as_1d_int(module.conv1d.dilation) + padding = _as_1d_int(module.conv1d.padding) + reset_gap = max(padding, dilation * (kernel_size - 1)) + + parts = [] + spans: list[tuple[int, int]] = [] + cursor = 0 + zero_gap = qkv.new_zeros((qkv.shape[0], reset_gap, qkv.shape[2])) if reset_gap > 0 else None + for idx, (start, end) in enumerate(ranges): + if idx > 0 and zero_gap is not None: + parts.append(zero_gap) + cursor += reset_gap + length = end - start + parts.append(qkv[:, start:end, :]) + spans.append((cursor, length)) + cursor += length + + qkv_padded = torch.cat(parts, dim=1) + groups = module.conv_dim_local_tp // module.cp_size + conv_out = F.conv1d( + input=qkv_padded.transpose(1, 2).contiguous(), + weight=conv1d_weight, + bias=conv1d_bias, + stride=module.conv1d.stride, + padding=module.conv1d.padding, + dilation=module.conv1d.dilation, + groups=groups, + ) + conv_out = module.act_fn(conv_out[..., : qkv_padded.shape[1]]).transpose(1, 2) + return torch.cat([conv_out[:, start : start + length, :] for start, length in spans], dim=1).contiguous() + + +def _fast_packed_causal_conv1d(module, qkv, conv1d_weight, conv1d_bias, cu_seqlens): + from fla.modules.convolution import causal_conv1d + + out, _ = causal_conv1d( + x=qkv.contiguous(), + weight=conv1d_weight.squeeze(1).contiguous(), + bias=conv1d_bias, + activation="silu", + cu_seqlens=cu_seqlens, + ) + return out.contiguous() + + +def apply_qwen_vl_packed_gdn_patch() -> None: + try: + from megatron.core.ssm import gated_delta_net as gdn_mod + from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx_range_push + except ImportError: + return + + GatedDeltaNet = gdn_mod.GatedDeltaNet + if getattr(GatedDeltaNet, "_slime_qwen_vl_packed_patch", False): + return + + original_forward = GatedDeltaNet.forward + + def patched_forward( + self, + hidden_states, + attention_mask, + inference_context=None, + packed_seq_params=None, + sequence_len_offset=None, + *, + inference_params=None, + **kwargs, + ): + if packed_seq_params is None: + return original_forward( + self, + hidden_states, + attention_mask, + inference_context=inference_context, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + inference_params=inference_params, + **kwargs, + ) + + inference_context = deprecate_inference_params(inference_context, inference_params) + if inference_context is not None: + raise NotImplementedError("GDN does not support inference for now.") + if self.config.deterministic_mode: + raise NotImplementedError("Packed GDN requires the FLA/FlashQLA kernel path.") + + seq_len, batch, _ = hidden_states.shape + seq_len = seq_len * self.sp_size * self.cp_size + cu_seqlens = _get_packed_cu_seqlens(packed_seq_params) + local_cu_seqlens = _local_cu_seqlens_for_kernel(self, cu_seqlens, seq_len) + should_time = _gdn_timing_should_sample(self) + timer = _CudaStepTimer(should_time) + packed_ranges = None + if should_time: + packed_ranges = _packed_ranges_from_cu_seqlens(local_cu_seqlens, seq_len) + + nvtx_range_push(suffix="in_proj") + qkvzba, _ = timer.run("in_proj", lambda: self.in_proj(hidden_states)) + nvtx_range_pop(suffix="in_proj") + + qkvzba = timer.run( + "cp2hp", + lambda: gdn_mod.tensor_a2a_cp2hp( + qkvzba, + seq_dim=0, + head_dim=-1, + cp_group=self.pg_collection.cp, + split_sections=[ + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + self.v_dim_local_tp, + self.num_value_heads // self.tp_size, + self.num_value_heads // self.tp_size, + ], + ), + ) + + qkvzba = qkvzba.transpose(0, 1) + qkv, gate, beta, alpha = torch.split( + qkvzba, + [ + (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // self.cp_size, + self.v_dim_local_tp // self.cp_size, + self.num_value_heads // self.tp_size // self.cp_size, + self.num_value_heads // self.tp_size // self.cp_size, + ], + dim=-1, + ) + gate = gate.reshape(batch, seq_len, -1, self.value_head_dim) + beta = beta.reshape(batch, seq_len, -1) + alpha = alpha.reshape(batch, seq_len, -1) + + nvtx_range_push(suffix="conv1d") + local_seq_len = qkv.shape[1] + qkv_channels_split_sections = [ + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + ] + conv1d_weight = gdn_mod.get_parameter_local_cp( + self.conv1d.weight, + dim=0, + cp_group=self.pg_collection.cp, + split_sections=qkv_channels_split_sections, + ) + conv1d_bias = ( + gdn_mod.get_parameter_local_cp( + self.conv1d.bias, + dim=0, + cp_group=self.pg_collection.cp, + split_sections=qkv_channels_split_sections, + ) + if self.conv_bias + else None + ) + use_fast_conv = _env_flag("SLIME_QWENVL_GDN_FAST_CONV") + if use_fast_conv and not getattr(self, "_slime_gdn_fast_conv_logged", False): + logger.info( + "Qwen VL packed GDN uses FLA causal_conv1d fast path with backend=%s.", + os.environ.get("SLIME_QWENVL_GDN_FAST_CONV_BACKEND", "triton"), + ) + self._slime_gdn_fast_conv_logged = True + use_batched_conv = (not use_fast_conv) and _env_flag("SLIME_QWENVL_GDN_BATCHED_CONV") + if use_batched_conv and not getattr(self, "_slime_gdn_batched_conv_logged", False): + logger.info("Qwen VL packed GDN uses batched safe torch causal conv.") + self._slime_gdn_batched_conv_logged = True + qkv = timer.run( + "conv1d", + lambda: ( + _fast_packed_causal_conv1d(self, qkv, conv1d_weight, conv1d_bias, local_cu_seqlens) + if use_fast_conv + else ( + _batched_safe_packed_causal_conv1d(self, qkv, conv1d_weight, conv1d_bias, local_cu_seqlens) + if use_batched_conv + else _safe_packed_causal_conv1d(self, qkv, conv1d_weight, conv1d_bias, local_cu_seqlens) + ) + ), + ) + qkv = qkv[:, :local_seq_len] + nvtx_range_pop(suffix="conv1d") + + query_key, value = torch.split( + qkv, + [2 * self.qk_dim_local_tp // self.cp_size, self.v_dim_local_tp // self.cp_size], + dim=-1, + ) + query_key = query_key.reshape(batch, seq_len, -1, self.key_head_dim) + value = value.reshape(batch, seq_len, -1, self.value_head_dim) + use_l2norm_in_kernel = self.use_qk_l2norm and _env_flag("SLIME_QWENVL_GDN_L2NORM_IN_KERNEL") + if use_l2norm_in_kernel and not getattr(self, "_slime_gdn_l2norm_in_kernel_logged", False): + logger.info("Qwen VL packed GDN uses kernel-side QK L2 norm.") + self._slime_gdn_l2norm_in_kernel_logged = True + if self.use_qk_l2norm and not use_l2norm_in_kernel: + query_key = timer.run("qk_l2norm", lambda: gdn_mod.l2norm(query_key.contiguous())) + + query, key = torch.split( + query_key, + [ + self.qk_dim_local_tp // self.key_head_dim // self.cp_size, + self.qk_dim_local_tp // self.key_head_dim // self.cp_size, + ], + dim=2, + ) + if self.num_value_heads // self.num_key_heads > 1: + query = query.repeat_interleave(self.num_value_heads // self.num_key_heads, dim=2) + key = key.repeat_interleave(self.num_value_heads // self.num_key_heads, dim=2) + + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + gate = gate.contiguous() + beta = beta.contiguous() + alpha = alpha.contiguous() + + nvtx_range_push(suffix="g_and_beta") + A_log_local_cp, dt_bias_local_cp = timer.run( + "g_params", + lambda: ( + gdn_mod.get_parameter_local_cp(self.A_log, dim=0, cp_group=self.pg_collection.cp), + gdn_mod.get_parameter_local_cp( + self.dt_bias, + dim=0, + cp_group=self.pg_collection.cp, + ), + ), + ) + g, beta = timer.run( + "g_and_beta", + lambda: ( + -A_log_local_cp.exp() * F.softplus(alpha.float() + dt_bias_local_cp), + beta.sigmoid(), + ), + ) + nvtx_range_pop(suffix="g_and_beta") + + nvtx_range_push(suffix="gated_delta_rule") + gated_delta_rule = _get_gated_delta_rule(self) + kernel_kwargs = {} + if getattr(self, "_slime_gdn_backend", "fla") == "flashqla": + kernel_kwargs["auto_cp"] = True + + core_attn_out, _ = timer.run( + "gated_delta_rule", + lambda: gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=use_l2norm_in_kernel, + cu_seqlens=local_cu_seqlens, + **kernel_kwargs, + ), + ) + nvtx_range_pop(suffix="gated_delta_rule") + + nvtx_range_push(suffix="gated_norm") + norm_out = timer.run("gated_norm", lambda: self._apply_gated_norm(core_attn_out, gate)) + nvtx_range_pop(suffix="gated_norm") + + norm_out = norm_out.reshape(batch, seq_len, -1) + norm_out = norm_out.transpose(0, 1).contiguous() + + norm_out = timer.run( + "hp2cp", + lambda: gdn_mod.tensor_a2a_hp2cp( + norm_out, + seq_dim=0, + head_dim=-1, + cp_group=self.pg_collection.cp, + ), + ) + + nvtx_range_push(suffix="out_proj") + out, out_bias = timer.run("out_proj", lambda: self.out_proj(norm_out)) + nvtx_range_pop(suffix="out_proj") + + if should_time: + logger.info( + "Qwen VL packed GDN timing: sample=%s batch=%s seq_len=%s local_seq_len=%s ranges=%s times_ms=%s", + getattr(self, "_slime_gdn_timing_count", 0), + batch, + seq_len, + local_seq_len, + _range_stats(packed_ranges or []), + {key: round(value, 3) for key, value in timer.summary().items()}, + ) + + return out, out_bias + + GatedDeltaNet.forward = patched_forward + GatedDeltaNet._slime_qwen_vl_packed_patch = True + logger.info("Patched Megatron-Core GatedDeltaNet packed forward for Qwen VL Bridge.") + + +apply_qwen_vl_packed_gdn_patch() diff --git a/slime/backends/megatron_utils/qwen_vl_text_language_fastpath.py b/slime/backends/megatron_utils/qwen_vl_text_language_fastpath.py new file mode 100644 index 0000000000..3c060cee9c --- /dev/null +++ b/slime/backends/megatron_utils/qwen_vl_text_language_fastpath.py @@ -0,0 +1,127 @@ +"""Opt-in text-only fastpath for Qwen3VL Bridge models. + +Qwen3VLModel owns the vision embedding + packed-sequence workaround. For +microbatches without multimodal tensors, Slime can already provide the normal +CP-local THD inputs used by the language model. This patch lets those batches +bypass the Qwen3VL wrapper while keeping the root model/DDP forward call. +""" + +import logging +import os + +import torch +from megatron.core import mpu + +logger = logging.getLogger(__name__) +_PATCHED = False + + +def _enabled() -> bool: + return os.environ.get("SLIME_QWENVL_TEXT_LANGUAGE_FASTPATH", "0").lower() in {"1", "true", "yes", "on"} + + +def _mrope_cp_thd_unsafe(kwargs: dict) -> bool: + packed_seq_params = kwargs.get("packed_seq_params") + if getattr(packed_seq_params, "qkv_format", None) != "thd": + return False + if kwargs.get("attention_mask") is not None: + return True + if kwargs.get("position_ids") is None: + return True + if os.environ.get("SLIME_QWENVL_TEXT_FASTPATH_LOCAL_MROPE", "1").lower() in {"1", "true", "yes", "on"}: + return False + try: + return mpu.get_context_parallel_world_size() > 1 + except Exception: + return True + + +def _has_tensor(value) -> bool: + if torch.is_tensor(value): + return value.numel() > 0 + if isinstance(value, dict): + return any(_has_tensor(v) for v in value.values()) + if isinstance(value, (list, tuple)): + return any(_has_tensor(v) for v in value) + return value is not None + + +def _has_mm_kwargs(kwargs: dict) -> bool: + for name in ( + "pixel_values", + "pixel_values_videos", + "image_grid_thw", + "video_grid_thw", + "image_input_mask", + "video_input_mask", + "cp_img_num", + "images_padded", + ): + if _has_tensor(kwargs.get(name)): + return True + return False + + +def apply_qwen_vl_text_language_fastpath_patch() -> None: + global _PATCHED + if _PATCHED: + return + try: + from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model import Qwen3VLModel + except Exception as exc: + logger.debug("QwenVL text language fastpath unavailable: %r", exc) + return + + original_forward = Qwen3VLModel.forward + + def _patched_forward(self, *args, **kwargs): + if not _enabled() or args or _has_mm_kwargs(kwargs): + return original_forward(self, *args, **kwargs) + + input_ids = kwargs.get("input_ids") + packed_seq_params = kwargs.get("packed_seq_params") + language_model = getattr(self, "language_model", None) + if language_model is None or input_ids is None or packed_seq_params is None: + return original_forward(self, *args, **kwargs) + + if getattr(packed_seq_params, "qkv_format", None) != "thd": + return original_forward(self, *args, **kwargs) + + if _mrope_cp_thd_unsafe(kwargs): + if not getattr(self, "_slime_qwenvl_text_fastpath_skip_logged", False): + logger.info( + "SLIME_QWENVL_TEXT_LANGUAGE_FASTPATH skipped for packed THD batch; " + "using Qwen3VL wrapper to preserve mRoPE/CP preprocessing" + ) + self._slime_qwenvl_text_fastpath_skip_logged = True + return original_forward(self, *args, **kwargs) + + if not getattr(self, "_slime_qwenvl_text_fastpath_logged", False): + logger.info("SLIME_QWENVL_TEXT_LANGUAGE_FASTPATH enabled: bypassing Qwen3VL wrapper for text-only THD batch") + self._slime_qwenvl_text_fastpath_logged = True + + rotary_pos_emb = getattr(language_model, "rotary_pos_emb", None) + if getattr(rotary_pos_emb, "is_thd_format", None) is not None: + rotary_pos_emb.is_thd_format = True + + return language_model( + input_ids=input_ids, + position_ids=kwargs.get("position_ids"), + attention_mask=kwargs.get("attention_mask"), + labels=kwargs.get("labels"), + inference_params=kwargs.get("inference_params"), + inference_context=kwargs.get("inference_context"), + packed_seq_params=packed_seq_params, + extra_block_kwargs=kwargs.get("extra_block_kwargs"), + runtime_gather_output=kwargs.get("runtime_gather_output"), + loss_mask=kwargs.get("loss_mask"), + visual_pos_masks=None, + deepstack_visual_embeds=None, + ) + + Qwen3VLModel.forward = _patched_forward + _PATCHED = True + logger.info("Installed QwenVL text language fastpath patch") + + +apply_qwen_vl_text_language_fastpath_patch() diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index f230c88f09..b63db5cad4 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -16,7 +16,6 @@ from slime.backends.sglang_utils.external import start_external_rollout_servers from slime.backends.sglang_utils.sglang_config import ModelConfig, ServerGroupConfig, SglangConfig -from slime.backends.sglang_utils.sglang_engine import SGLangEngine from slime.rollout.base_types import call_rollout_fn from slime.utils import logging_utils from slime.utils.dp_schedule import build_dp_schedule @@ -102,6 +101,20 @@ def _tensorize_rollout_data_for_training(rollout_data: dict[str, Any]) -> None: ) +def _has_multimodal_train_inputs(value: Any) -> bool: + if value is None: + return False + if isinstance(value, dict): + return any(_has_multimodal_train_inputs(v) for v in value.values()) + if isinstance(value, (list, tuple)): + return any(_has_multimodal_train_inputs(v) for v in value) + if isinstance(value, np.ndarray): + return value.size > 0 + if torch.is_tensor(value): + return value.numel() > 0 + return True + + @dataclasses.dataclass class ServerGroup: """A group of homogeneous SGLang engines with the same configuration. @@ -165,6 +178,8 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis rollout_num_gpus_per_engine=self.args.rollout_num_gpus_per_engine, ) + from slime.backends.sglang_utils.sglang_engine import SGLangEngine + RolloutRayActor = ray.remote(SGLangEngine) rollout_engines = [] @@ -841,6 +856,11 @@ def _split_train_data_by_dp(self, data): dp_size = self.train_parallel_config["dp_size"] total_lengths = [len(t) for t in data["tokens"]] data["total_lengths"] = total_lengths + sample_is_multimodal = None + if getattr(self.args, "multimodal_aware_packing", "off") != "off" and "multimodal_train_inputs" in data: + sample_is_multimodal = [ + _has_multimodal_train_inputs(mm_inputs) for mm_inputs in data["multimodal_train_inputs"] + ] partitions, micro_batch_indices, num_microbatches, global_batch_sizes = build_dp_schedule( self.args, @@ -848,7 +868,26 @@ def _split_train_data_by_dp(self, data): total_lengths, global_batch_size=self.args.global_batch_size, rollout_indices=data["rollout_ids"], + sample_is_multimodal=sample_is_multimodal, ) + if getattr(self.args, "global_batch_tokens", None) is not None: + kept_samples = sum(len(partition) for partition in partitions) + logger.info( + "token-based global batch: steps=%s target=%s tokens/step global_batch_sizes=%s kept_samples=%s/%s", + len(global_batch_sizes), + self.args.global_batch_tokens, + global_batch_sizes, + kept_samples, + len(total_lengths), + ) + if sample_is_multimodal is not None: + logger.info( + "multimodal-aware packing: mode=%s multimodal_samples=%s/%s num_microbatches=%s", + self.args.multimodal_aware_packing, + sum(sample_is_multimodal), + len(sample_is_multimodal), + num_microbatches, + ) # Package per-rank rollout_data rollout_data_refs = [] diff --git a/slime/rollout/sft_rollout.py b/slime/rollout/sft_rollout.py index 4407463d23..b728cb8ed5 100644 --- a/slime/rollout/sft_rollout.py +++ b/slime/rollout/sft_rollout.py @@ -1,7 +1,15 @@ import logging +import json +import os +from copy import deepcopy +from typing import Any + +import numpy as np +import torch from slime.utils.mask_utils import MultiTurnLossMaskGenerator -from slime.utils.processing_utils import load_processor, load_tokenizer +from slime.utils.processing_utils import build_processor_kwargs, load_processor, load_tokenizer +from slime.utils.types import Sample __all__ = ["generate_rollout"] @@ -14,6 +22,165 @@ SAMPLE_PRINTED = False +class OverlongSFTSampleError(ValueError): + def __init__(self, length: int, max_context_len: int): + super().__init__(f"SFT sample length {length} exceeds --rollout-max-context-len {max_context_len}.") + self.length = length + self.max_context_len = max_context_len + + +def _env_flag(name: str, default: str = "0") -> bool: + return os.environ.get(name, default).lower() in {"1", "true", "yes", "on"} + + +def _to_list_ids(input_ids: Any) -> list[int]: + if isinstance(input_ids, torch.Tensor): + return input_ids.tolist() + if isinstance(input_ids, np.ndarray): + return input_ids.tolist() + return list(input_ids) + + +def _first_sequence_ids(input_ids: Any) -> list[int]: + if isinstance(input_ids, torch.Tensor): + return _to_list_ids(input_ids[0] if input_ids.ndim > 1 else input_ids) + if isinstance(input_ids, np.ndarray): + return _to_list_ids(input_ids[0] if input_ids.ndim > 1 else input_ids) + if input_ids and isinstance(input_ids[0], (list, tuple)): + return _to_list_ids(input_ids[0]) + return _to_list_ids(input_ids) + + +def _to_tensor_or_value(value: Any) -> Any: + if isinstance(value, np.ndarray): + return torch.from_numpy(value) + return value + + +def _messages_from_sample(sample: Sample) -> list[dict[str, Any]]: + prompt = deepcopy(sample.prompt) + + if isinstance(prompt, str): + if sample.label is None: + raise ValueError("SFT sample has a string prompt but no label to form the assistant response.") + messages = [ + {"role": "user", "content": prompt}, + {"role": "assistant", "content": sample.label}, + ] + elif isinstance(prompt, list): + messages = prompt + has_assistant = any(message.get("role") == "assistant" for message in messages) + if not has_assistant: + if sample.label is None: + raise ValueError("SFT sample has no assistant turn and no label.") + messages = messages + [{"role": "assistant", "content": sample.label}] + else: + raise TypeError(f"Unsupported SFT prompt type: {type(prompt)}") + + return _normalize_tool_call_arguments(messages) + + +def _normalize_tool_call_arguments(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Normalize OpenAI tool-call arguments to the mapping expected by Qwen templates.""" + for message in messages: + tool_calls = message.get("tool_calls") + if not isinstance(tool_calls, list): + continue + for tool_call in tool_calls: + if not isinstance(tool_call, dict): + continue + function = tool_call.get("function") + if not isinstance(function, dict): + continue + arguments = function.get("arguments") + if isinstance(arguments, str): + try: + parsed_arguments = json.loads(arguments) + except json.JSONDecodeError: + continue + if isinstance(parsed_arguments, dict): + function["arguments"] = parsed_arguments + return messages + + +def _get_chat_template_kwargs(args) -> dict: + kwargs = getattr(args, "apply_chat_template_kwargs", None) + if kwargs is None: + return {} + if isinstance(kwargs, str): + if not kwargs.strip(): + return {} + kwargs = json.loads(kwargs) + if not isinstance(kwargs, dict): + raise TypeError(f"apply_chat_template_kwargs must be a dict or JSON string, got {type(kwargs)}") + return dict(kwargs) + + +def _render_messages( + messages: list[dict[str, Any]], tools: list[dict] | None, chat_template_kwargs: dict | None = None +) -> str: + return TOKENIZER.apply_chat_template( + messages, + tokenize=False, + tools=tools, + return_dict=False, + **(chat_template_kwargs or {}), + ) + + +def _build_sft_sample(args, sample: Sample) -> Sample: + messages = _messages_from_sample(sample) + tools = sample.metadata.get("tools", None) if sample.metadata else None + + if args.multimodal_keys is not None: + if PROCESSOR is None: + raise RuntimeError("--multimodal-keys is set, but no HuggingFace processor was loaded.") + + rendered_text = _render_messages(messages, tools, _get_chat_template_kwargs(args)) + processor_output = PROCESSOR( + text=[rendered_text], + **build_processor_kwargs(sample.multimodal_inputs or {}), + ) + input_ids = _first_sequence_ids(processor_output["input_ids"]) + token_ids, loss_mask = MASK_GENERATOR.get_loss_mask_with_multimodal_alignment( + messages, + input_ids, + tools=tools, + rendered_text=rendered_text, + image_grid_thw=processor_output.get("image_grid_thw"), + image_token=getattr(PROCESSOR, "image_token", "<|image_pad|>"), + image_merge_size=getattr(PROCESSOR.image_processor, "merge_size", 2), + ) + sample.multimodal_train_inputs = { + key: _to_tensor_or_value(value) + for key, value in processor_output.items() + if key not in ("input_ids", "attention_mask") + } or None + else: + token_ids, loss_mask = MASK_GENERATOR.get_loss_mask(messages, tools=tools) + sample.multimodal_train_inputs = None + + if len(token_ids) != len(loss_mask): + raise ValueError( + f"SFT rollout produced mismatched token_ids/loss_mask lengths: {len(token_ids)=}, {len(loss_mask)=}" + ) + + max_context_len = getattr(args, "rollout_max_context_len", None) + if max_context_len is not None and len(token_ids) > max_context_len: + raise OverlongSFTSampleError(len(token_ids), max_context_len) + + response_length = MASK_GENERATOR.get_response_lengths([loss_mask])[0] + if response_length <= 0: + raise ValueError(f"SFT sample has no supervised assistant tokens: {messages=}") + + sample.tokens = token_ids + sample.response_length = response_length + sample.reward = 0.0 + sample.loss_mask = loss_mask[-response_length:] + sample.status = Sample.Status.COMPLETED + return sample + + def generate_rollout(args, rollout_id, data_buffer, evaluation=False): """An example to implement the generate_rollout function for an rule based rm rollout generation. @@ -37,32 +204,74 @@ def generate_rollout(args, rollout_id, data_buffer, evaluation=False): PROCESSOR = load_processor(args.hf_checkpoint, trust_remote_code=True) if MASK_GENERATOR is None: - MASK_GENERATOR = MultiTurnLossMaskGenerator(TOKENIZER, tokenizer_type=args.loss_mask_type) + MASK_GENERATOR = MultiTurnLossMaskGenerator( + TOKENIZER, + tokenizer_type=args.loss_mask_type, + chat_template_kwargs=_get_chat_template_kwargs(args), + ) - samples = data_buffer.get_samples(args.rollout_batch_size) + skip_overlong = _env_flag("SLIME_SFT_SKIP_OVERLONG", "0") + max_fetch_factor = max(1, int(os.environ.get("SLIME_SFT_SKIP_OVERLONG_MAX_FACTOR", "8"))) + max_attempts = args.rollout_batch_size * max_fetch_factor + sample_groups = [] + attempted_samples = 0 + skipped_overlong = 0 + pending_groups = data_buffer.get_samples(args.rollout_batch_size) - for i, sample in enumerate(samples): - (sample,) = sample - messages = sample.prompt - tools = sample.metadata.get("tools", None) + while pending_groups and len(sample_groups) < args.rollout_batch_size: + for group in pending_groups: + (sample,) = group + attempted_samples += 1 + try: + _build_sft_sample(args, sample) + except OverlongSFTSampleError as exc: + if not skip_overlong: + raise + skipped_overlong += 1 + if skipped_overlong <= 8: + logger.warning( + "Skipping overlong SFT sample: length=%s max_context_len=%s", + exc.length, + exc.max_context_len, + ) + continue - token_ids, loss_mask = MASK_GENERATOR.get_loss_mask(messages, tools=tools) - if len(token_ids) != len(loss_mask): - raise ValueError( - f"SFT rollout produced mismatched token_ids/loss_mask lengths: {len(token_ids)=}, {len(loss_mask)=}" - ) + sample_groups.append([sample]) + if len(sample_groups) >= args.rollout_batch_size: + break + + remaining = args.rollout_batch_size - len(sample_groups) + if not skip_overlong or remaining <= 0: + break + if attempted_samples >= max_attempts: + break + pending_groups = data_buffer.get_samples(remaining) - response_length = MASK_GENERATOR.get_response_lengths([loss_mask])[0] + if len(sample_groups) < args.rollout_batch_size: + raise RuntimeError( + f"Only built {len(sample_groups)}/{args.rollout_batch_size} SFT samples after " + f"{attempted_samples} attempts; skipped_overlong={skipped_overlong}. " + "Increase SLIME_SFT_SKIP_OVERLONG_MAX_FACTOR or clean the dataset." + ) - sample.tokens = token_ids - sample.response_length = response_length - sample.reward = 0 - sample.loss_mask = loss_mask[-response_length:] + if skipped_overlong: + logger.warning( + "Skipped %s overlong SFT samples while building rollout batch of %s " + "(attempted_samples=%s, max_attempts=%s).", + skipped_overlong, + args.rollout_batch_size, + attempted_samples, + max_attempts, + ) - if i == 0 and not SAMPLE_PRINTED: - logger.info( - f"sft_rollout::generate_rollout example data: {sample=} (raw){messages=} (raw){token_ids=} (raw){loss_mask=} {response_length=}" - ) - SAMPLE_PRINTED = True + if not SAMPLE_PRINTED: + (sample,) = sample_groups[0] + mm_keys = sorted(sample.multimodal_train_inputs.keys()) if sample.multimodal_train_inputs else [] + logger.info( + "sft_rollout example: " + f"tokens={len(sample.tokens)}, response_length={sample.response_length}, " + f"loss_tokens={sum(sample.loss_mask or [])}, multimodal_train_input_keys={mm_keys}" + ) + SAMPLE_PRINTED = True - return samples + return sample_groups diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index b29d5cf0cf..897d738714 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -7,9 +7,6 @@ import yaml -from slime.backends.sglang_utils.arguments import sglang_parse_args -from slime.backends.sglang_utils.arguments import validate_args as sglang_validate_args -from slime.backends.sglang_utils.external import apply_external_engine_info_to_args from slime.utils.eval_config import EvalDatasetConfig, build_eval_dataset_configs, ensure_dataset_list from slime.utils.logging_utils import configure_logger @@ -106,6 +103,31 @@ def add_cluster_arguments(parser): def add_train_arguments(parser): # --train-backend is parsed early in _pre_parse_mode() and merged later. + parser.add_argument( + "--qkv-format", + type=str, + choices=["thd", "bshd"], + default="thd", + help="The qkv layout for Megatron backend.", + ) + # Vendored Megatron exposes these TransformerConfig fields but does not + # register CLI flags for them, and Slime ignores unknown Megatron args. + reset_arg( + parser, + "--fp8-dot-product-attention", + action="store_true", + default=False, + dest="fp8_dot_product_attention", + help="Use TransformerEngine FP8 Dot Product Attention when FP8 is enabled.", + ) + reset_arg( + parser, + "--fp8-multi-head-attention", + action="store_true", + default=False, + dest="fp8_multi_head_attention", + help="Use TransformerEngine FP8 Multi Head Attention when FP8 is enabled.", + ) parser.add_argument( "--qwen-gdn-backend", type=str, @@ -292,6 +314,15 @@ def add_train_arguments(parser): --freeze-params-name-list linear_fc1 """, ) + parser.add_argument( + "--vision-dp-when-cp", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "For QwenVL Bridge models under context parallelism, shard the vision tower work by CP rank " + "and all-gather the resulting vision embeddings before the language model." + ), + ) parser.add_argument( "--allgather-cp", action="store_true", @@ -700,6 +731,16 @@ def add_data_arguments(parser): "`rollout_batch_size * n_samples_per_prompt // num_steps_per_rollout`." ), ) + parser.add_argument( + "--global-batch-tokens", + type=int, + default=None, + help=( + "Enable token-based global batches. When set, each training step holds a near-precise " + "token budget instead of a fixed rollout/sample count. Requires --use-dynamic-batch-size; " + "incompatible with --use-dynamic-global-batch-size and --balance-data." + ), + ) # mbs for the training, will be ignored if `use_dynamic_batch_size` is set. reset_arg(parser, "--micro-batch-size", type=int, default=1) parser.add_argument( @@ -751,6 +792,29 @@ def add_data_arguments(parser): "`max_response_len // cp_size` instead of `max_response_len`." ), ) + parser.add_argument( + "--packing-safety-margin", + type=float, + default=1.0, + help=( + "Multiplier applied to max_tokens_per_gpu when planning dynamic micro-batches. " + "Values below 1.0 leave headroom for activation/padding spikes while validation still " + "uses the real max_tokens_per_gpu cap." + ), + ) + parser.add_argument( + "--multimodal-aware-packing", + type=str, + choices=["off", "separate", "separate_raw"], + default="off", + help=( + "Optional dynamic micro-batch packing strategy for multimodal SFT. " + "'off' preserves the standard token-sum first-fit packing. " + "'separate' avoids mixing true multimodal samples with text-only samples and " + "uses padded unsplit cost for multimodal bins. " + "'separate_raw' also avoids mixing but keeps the standard token-sum bin cost." + ), + ) parser.add_argument( "--log-probs-max-tokens-per-gpu", type=int, @@ -1556,6 +1620,8 @@ def parse_args(add_custom_arguments=None): # Skipped when sglang servers are not needed. sglang_ns = None if not skip_sglang: + from slime.backends.sglang_utils.arguments import sglang_parse_args + sglang_ns = sglang_parse_args() # Phase 2: Parse megatron + slime args. @@ -1584,6 +1650,8 @@ def parse_args(add_custom_arguments=None): megatron_validate_args(args) if not args.debug_train_only: + from slime.backends.sglang_utils.arguments import validate_args as sglang_validate_args + sglang_validate_args(args) return args @@ -1816,12 +1884,19 @@ def slime_validate_args(args): if args.use_dynamic_batch_size: assert args.max_tokens_per_gpu is not None, "max_tokens_per_gpu must be set when use_dynamic_batch_size is set" + assert args.packing_safety_margin > 0, "packing_safety_margin must be positive" if args.log_probs_max_tokens_per_gpu is None: args.log_probs_max_tokens_per_gpu = args.max_tokens_per_gpu - if getattr(args, "balance_by_flops", False): - assert args.use_dynamic_batch_size, "--balance-by-flops requires --use-dynamic-batch-size" - args.balance_data = True + if args.global_batch_tokens is not None: + assert args.global_batch_tokens > 0, "global_batch_tokens must be positive" + assert ( + args.use_dynamic_batch_size + ), "--global-batch-tokens requires --use-dynamic-batch-size (token-based microbatch packing)" + assert ( + not getattr(args, "use_dynamic_global_batch_size", False) + ), "--global-batch-tokens is incompatible with --use-dynamic-global-batch-size" + assert not args.balance_data, "--global-batch-tokens requires a round-robin DP split; disable --balance-data" if args.eps_clip_high is None: args.eps_clip_high = args.eps_clip diff --git a/slime/utils/data.py b/slime/utils/data.py index 0d26b6dda5..cbde4d1485 100644 --- a/slime/utils/data.py +++ b/slime/utils/data.py @@ -4,6 +4,7 @@ import os import random import re +from collections import defaultdict import numpy as np import ray @@ -140,62 +141,68 @@ def _build_messages(data: dict, prompt_key: str, as_conversation: bool, multimod if multimodal_keys: # Build mapping: placeholder -> (MultimodalType, content_list) multimodals = {} + remain_data = defaultdict(int) for type_name, data_key in multimodal_keys.items(): mt = MultimodalTypes.get(type_name) - if mt: - multimodal_data = data.get(data_key) - if multimodal_data is not None: - multimodals[mt.placeholder] = (mt, list(multimodal_data)) - - pattern = "(" + "|".join(re.escape(p) for p in multimodals.keys()) + ")" - - for message in prompt: - if isinstance(message["content"], str): - content_list = [] - for segment in re.split(pattern, message["content"]): - if not segment: - continue - if segment in multimodals: - mt, content = multimodals[segment] - assert len(content) > 0, ( - f"Not enough {mt.name} data: more '{mt.placeholder}' placeholders in prompt " - f"than {mt.name}s provided in data" - ) - item = content.pop(0) - # Support rich image config from https://github.com/QwenLM/Qwen3-VL/blob/main/README.md - # "images": [{"type": "image", "image": "path/to/img/01.jpeg", "max_pixels": 50176, "min_pixels": 50176}, {...}] - if isinstance(item, dict): - content_list.append(item) - # "images": ["path/to/img/01.jpeg", "url", "base64enc"] - else: - content_list.append({"type": mt.name, mt.name: item}) - else: - content_list.append({"type": "text", "text": segment}) - message["content"] = content_list - - elif isinstance(message["content"], list): - # TODO: handle more general cases. where message['content'] is a dict and contains multiple types of content. - # e.g. - # "content": [ - # { - # "type": "image", - # "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg", - # }, - # {"type": "text", "text": "Describe this image."}, - # ], - logger.warning("message['content'] is a list of dicts, no processing will be done.") - continue + if mt is None: + raise ValueError(f"Unsupported multimodal type: {type_name}") + + raw_multimodal_data = data.get(data_key) + if raw_multimodal_data is None: + multimodal_data = [] + elif isinstance(raw_multimodal_data, (str, bytes)): + multimodal_data = [raw_multimodal_data] else: - raise ValueError( - f"Unsupported content type: {type(message['content'])}, expected str or list of dicts" + multimodal_data = list(raw_multimodal_data) + + multimodals[mt.placeholder] = (mt, multimodal_data) + remain_data[mt.name] += len(multimodal_data) + + if multimodals: + pattern = "(" + "|".join(re.escape(p) for p in multimodals.keys()) + ")" + built_prompt = [] + + for message in prompt: + if isinstance(message["content"], str): + content_list = [] + for segment in re.split(pattern, message["content"]): + if not segment: + continue + if segment in multimodals: + mt, content = multimodals[segment] + remain_data[mt.name] -= 1 + if remain_data[mt.name] < 0: + logger.warning( + "The number of placeholder %s in prompt is more than data number.", segment + ) + if content: + content_list.append({"type": mt.name, mt.name: content.pop(0)}) + else: + content_list.append({"type": "text", "text": segment}) + + built_message = dict(message) + built_message["content"] = content_list + built_prompt.append(built_message) + + elif isinstance(message["content"], list): + for item in message["content"]: + item_type = item.get("type") if isinstance(item, dict) else None + if item_type in remain_data: + remain_data[item_type] -= 1 + built_prompt.append(message) + else: + raise ValueError( + f"Unsupported content type: {type(message['content'])}, expected str or list of dicts" + ) + + prompt = built_prompt + + if any(v > 0 for v in remain_data.values()): + raise RuntimeError( + f"placeholder lost! The number of remain mutimodal data is {remain_data}. " + "Please check your dataset prompt." ) - for placeholder, (mt, remaining) in multimodals.items(): - assert len(remaining) == 0, ( - f"Multimodal data count mismatch: {len(remaining)} more {mt.name}(s)" - f"than '{placeholder}' placeholders in prompt" - ) - return prompt diff --git a/slime/utils/dp_schedule.py b/slime/utils/dp_schedule.py index 1735fad3b4..622a70af4b 100644 --- a/slime/utils/dp_schedule.py +++ b/slime/utils/dp_schedule.py @@ -59,26 +59,144 @@ def _pack_step_into_mbs( use_dynamic_batch_size: bool, max_per_bin: int | None, micro_batch_size: int | None, - balance_by_flops: bool = False, + multimodal_aware_packing: str = "off", + step_is_multimodal: list[bool] | None = None, ) -> list[list[int]]: """Group a step's samples into mbs. Returns ``mbs[k]`` = local indices into ``step_lengths``.""" if use_dynamic_batch_size: assert max_per_bin is not None - if balance_by_flops: - total_tokens = sum(step_lengths) - num_mbs = max(1, (total_tokens + max_per_bin - 1) // max_per_bin) - if num_mbs >= len(step_lengths): - return [[i] for i in range(len(step_lengths))] - workloads = _calculate_workloads(step_lengths, args) - # FLOPs balancing does not enforce the token cap per mbs. A - # partition can exceed max_per_bin and may OOM if the cap is tight. - return get_seqlen_balanced_partitions(workloads, num_mbs, equal_size=False) + if ( + multimodal_aware_packing in {"separate", "separate_raw"} + and step_is_multimodal is not None + and any(step_is_multimodal) + ): + return _first_fit_pack_separate_multimodal( + step_lengths, + step_is_multimodal, + max_per_bin, + padded_multimodal_cost=multimodal_aware_packing == "separate", + ) return first_fit_pack(step_lengths, max_per_bin) assert micro_batch_size is not None n = len(step_lengths) return [list(range(i, min(i + micro_batch_size, n))) for i in range(0, n, micro_batch_size)] +def _first_fit_pack_separate_multimodal( + step_lengths: list[int], + step_is_multimodal: list[bool], + max_tokens_per_bin: int, + *, + padded_multimodal_cost: bool, +) -> list[list[int]]: + """First-fit packing that does not mix true multimodal and text-only samples. + + ``separate_raw`` uses the usual sum of token lengths for both text-only and + multimodal bins while still avoiding text/mm mixing. ``separate`` uses + ``len(bin) * max(seq_len in bin)`` for multimodal bins because the QwenVL + unsplit path pads samples in the same microbatch to the longest sequence + before entering the wrapper. + """ + assert len(step_lengths) == len(step_is_multimodal) + + bins: list[list[int]] = [] + bin_is_multimodal: list[bool] = [] + bin_sums: list[int] = [] + bin_maxes: list[int] = [] + + for idx, length in enumerate(step_lengths): + is_mm = step_is_multimodal[idx] + for j, bin_ in enumerate(bins): + if bin_is_multimodal[j] != is_mm: + continue + if is_mm and padded_multimodal_cost: + new_max = max(bin_maxes[j], length) + new_cost = new_max * (len(bin_) + 1) + else: + new_cost = bin_sums[j] + length + if new_cost <= max_tokens_per_bin: + bin_.append(idx) + bin_sums[j] += length + bin_maxes[j] = max(bin_maxes[j], length) + break + else: + bins.append([idx]) + bin_is_multimodal.append(is_mm) + bin_sums.append(length) + bin_maxes.append(length) + + return bins + + +def _group_samples_by_rollout(rollout_indices: list[int]) -> tuple[list[int], dict[int, list[int]]]: + rollout_id_to_samples: dict[int, list[int]] = {} + for sample_pos, rid in enumerate(rollout_indices): + rollout_id_to_samples.setdefault(rid, []).append(sample_pos) + return list(rollout_id_to_samples.keys()), rollout_id_to_samples + + +def _fixed_rollout_steps(rollout_ids: list[int], global_batch_size: int) -> list[list[int]]: + num_steps = len(rollout_ids) // global_batch_size + assert num_steps >= 1, ( + f"num_rollouts ({len(rollout_ids)}) < global_batch_size ({global_batch_size}); " + f"need at least one rollout per step." + ) + return [rollout_ids[i * global_batch_size : (i + 1) * global_batch_size] for i in range(num_steps)] + + +def _token_budget_rollout_steps( + rollout_ids: list[int], + rollout_id_to_samples: dict[int, list[int]], + total_lengths: list[int], + target_tokens: int, + dp_size: int, +) -> list[list[int]]: + if target_tokens <= 0: + raise ValueError(f"target_tokens must be positive, got {target_tokens}") + + steps: list[list[int]] = [] + start = 0 + while start < len(rollout_ids): + chosen = len(rollout_ids) + prev = None + c = start + 1 + while c <= len(rollout_ids): + step_rollouts = rollout_ids[start:c] + sample_count = sum(len(rollout_id_to_samples[rid]) for rid in step_rollouts) + token_count = sum( + total_lengths[i] + for rid in step_rollouts + for i in rollout_id_to_samples[rid] + ) + dp_aligned = sample_count >= dp_size and sample_count % dp_size == 0 + if dp_aligned and token_count >= target_tokens: + if prev is not None: + prev_rollouts = rollout_ids[start:prev] + prev_tokens = sum( + total_lengths[i] + for rid in prev_rollouts + for i in rollout_id_to_samples[rid] + ) + under_gap = target_tokens - prev_tokens + over_gap = token_count - target_tokens + chosen = prev if under_gap < over_gap else c + else: + chosen = c + break + if dp_aligned: + prev = c + c += 1 + + step_rollouts = rollout_ids[start:chosen] + sample_count = sum(len(rollout_id_to_samples[rid]) for rid in step_rollouts) + if sample_count < dp_size or sample_count % dp_size != 0: + break + steps.append(step_rollouts) + start = chosen + + return steps + + def build_dp_schedule( args: Any, train_parallel_config: dict, @@ -86,6 +204,7 @@ def build_dp_schedule( *, global_batch_size: int, rollout_indices: list[int], + sample_is_multimodal: list[bool] | None = None, ) -> tuple[list[list[int]], list[list[list[int]]], list[int], list[int]]: """Compute the per-rank DP partition and micro-batch schedule. @@ -103,6 +222,10 @@ def build_dp_schedule( samples don't fit are dropped. rollout_indices: rollout id for each sample (``samples[i].index``). Samples sharing the same id are kept together in one step. + sample_is_multimodal: optional boolean per sample. When + ``args.multimodal_aware_packing == "separate"``, dynamic packing + avoids mixing multimodal and text-only samples and estimates + multimodal mbs cost by padded unsplit size. Returns: ``(partitions, micro_batch_indices, num_microbatches, global_batch_sizes)``. @@ -113,11 +236,18 @@ def build_dp_schedule( cp_size = train_parallel_config["cp_size"] vpp_size = train_parallel_config["vpp_size"] mb_group = train_parallel_config["microbatch_group_size_per_vp_stage"] + if sample_is_multimodal is not None: + assert len(sample_is_multimodal) == len(total_lengths), ( + f"sample_is_multimodal length {len(sample_is_multimodal)} does not match " + f"total_lengths length {len(total_lengths)}" + ) max_per_bin = None if args.use_dynamic_batch_size: assert args.max_tokens_per_gpu is not None - max_per_bin = args.max_tokens_per_gpu * cp_size + max_tokens = args.max_tokens_per_gpu * cp_size + packing_safety_margin = getattr(args, "packing_safety_margin", 1.0) + max_per_bin = max(1, int(max_tokens * packing_safety_margin)) # mbs count per step must be divisible by (dp_size * mb_group_for_vpp) so # every rank ends up with the same num_mbs and (for VPP) the per-rank mbs @@ -127,27 +257,32 @@ def build_dp_schedule( # Group samples by rollout id (preserve first-occurrence order). All # samples from one rollout stay in a single step so the per-rollout loss # reducer is well-defined. - rollout_id_to_samples: dict[int, list[int]] = {} - for sample_pos, rid in enumerate(rollout_indices): - rollout_id_to_samples.setdefault(rid, []).append(sample_pos) - rollout_ids = list(rollout_id_to_samples.keys()) - - num_steps = len(rollout_ids) // global_batch_size - assert num_steps >= 1, ( - f"num_rollouts ({len(rollout_ids)}) < global_batch_size ({global_batch_size}); " - f"need at least one rollout per step." - ) + rollout_ids, rollout_id_to_samples = _group_samples_by_rollout(rollout_indices) + if getattr(args, "global_batch_tokens", None) is not None: + step_rollout_groups = _token_budget_rollout_steps( + rollout_ids, + rollout_id_to_samples, + total_lengths, + args.global_batch_tokens, + dp_size, + ) + assert step_rollout_groups, ( + f"num_rollouts ({len(rollout_ids)}) cannot form a token-based global batch " + f"with dp_size {dp_size} and global_batch_tokens {args.global_batch_tokens}" + ) + else: + step_rollout_groups = _fixed_rollout_steps(rollout_ids, global_batch_size) partitions: list[list[int]] = [[] for _ in range(dp_size)] micro_batch_indices: list[list[list[int]]] = [[] for _ in range(dp_size)] num_microbatches: list[int] = [] global_batch_sizes: list[int] = [] - for step_i in range(num_steps): - step_rollouts = rollout_ids[step_i * global_batch_size : (step_i + 1) * global_batch_size] + for step_i, step_rollouts in enumerate(step_rollout_groups): sample_indices = [pos for rid in step_rollouts for pos in rollout_id_to_samples[rid]] step_lengths = [total_lengths[i] for i in sample_indices] - global_batch_sizes.append(global_batch_size) + step_is_multimodal = [sample_is_multimodal[i] for i in sample_indices] if sample_is_multimodal is not None else None + global_batch_sizes.append(len(step_rollouts)) assert len(sample_indices) >= dp_size, ( f"step {step_i}: {len(sample_indices)} samples < dp_size {dp_size}; " f"each step needs at least one sample per rank." @@ -161,7 +296,8 @@ def build_dp_schedule( use_dynamic_batch_size=args.use_dynamic_batch_size, max_per_bin=max_per_bin, micro_batch_size=getattr(args, "micro_batch_size", None), - balance_by_flops=args.balance_by_flops, + multimodal_aware_packing=getattr(args, "multimodal_aware_packing", "off"), + step_is_multimodal=step_is_multimodal, ) # 2. Align mbs count to a multiple of ``align_to``. diff --git a/slime/utils/mask_utils.py b/slime/utils/mask_utils.py index d298946106..ef3ad52907 100644 --- a/slime/utils/mask_utils.py +++ b/slime/utils/mask_utils.py @@ -1,3 +1,5 @@ +import json + from transformers import AutoTokenizer @@ -6,10 +8,23 @@ def get_response_lengths(loss_masks: list[list[int]]) -> list[int]: return [len(mask[mask.index(1) :]) if 1 in mask else 0 for mask in loss_masks] +def _normalize_chat_template_kwargs(chat_template_kwargs) -> dict: + if chat_template_kwargs is None: + return {} + if isinstance(chat_template_kwargs, str): + if not chat_template_kwargs.strip(): + return {} + chat_template_kwargs = json.loads(chat_template_kwargs) + if not isinstance(chat_template_kwargs, dict): + raise TypeError(f"chat_template_kwargs must be a dict or JSON string, got {type(chat_template_kwargs)}") + return dict(chat_template_kwargs) + + class MultiTurnLossMaskGenerator: - def __init__(self, tokenizer: AutoTokenizer, tokenizer_type: str = "qwen"): + def __init__(self, tokenizer: AutoTokenizer, tokenizer_type: str = "qwen", chat_template_kwargs=None): self.tokenizer = tokenizer self.tokenizer_type = tokenizer_type + self.chat_template_kwargs = _normalize_chat_template_kwargs(chat_template_kwargs) self.system_message_length = 0 self.gen_token_length = 0 if self.tokenizer_type in ("qwen", "qwen3"): @@ -127,116 +142,82 @@ def gen_multi_turn_loss_mask_qwen3( def gen_multi_turn_loss_mask_qwen3_5( self, messages: list[dict], tools: list[dict] = None ) -> tuple[list[int], list[int]]: - rendered_text = self.tokenizer.apply_chat_template(messages, tokenize=False, tools=tools, return_dict=False) - tokenized = self.tokenizer(rendered_text, add_special_tokens=False, return_offsets_mapping=True) - token_ids = tokenized["input_ids"] - offset_mapping = tokenized.get("offset_mapping") - - if offset_mapping is None: - raise ValueError( - "Qwen3.5 loss mask generation requires a fast tokenizer " "with `return_offsets_mapping` support." - ) + rendered_text = self.tokenizer.apply_chat_template( + messages, + tokenize=False, + tools=tools, + return_dict=False, + **self.chat_template_kwargs, + ) + token_ids, loss_mask = self.get_loss_mask_from_rendered_qwen3_5(messages, rendered_text) expected_token_ids = self.tokenizer.apply_chat_template( - messages, tokenize=True, tools=tools, return_dict=False + messages, + tokenize=True, + tools=tools, + return_dict=False, + **self.chat_template_kwargs, ) if token_ids != expected_token_ids: raise ValueError( - "Qwen3.5 rendered text tokenization does not match " + "Qwen3.5/Qwen3.6 rendered text tokenization does not match " "`apply_chat_template(..., tokenize=True)` output." ) - assistant_header = "<|im_start|>assistant\n" - think_prefix = "\n" - end_marker = "<|im_end|>" - - char_mask = [0] * len(rendered_text) - cursor = 0 - - for message in messages: - if message["role"] != "assistant": - continue - - header_pos = rendered_text.find(assistant_header, cursor) - if header_pos < 0: - raise ValueError("Failed to locate assistant message in rendered Qwen3.5 chat template output.") - - content_start = header_pos + len(assistant_header) - end_pos = rendered_text.find(end_marker, content_start) - if end_pos < 0: - raise ValueError("Failed to locate <|im_end|> for assistant message in rendered text.") - - span_end = end_pos + len(end_marker) - if span_end < len(rendered_text) and rendered_text[span_end] == "\n": - span_end += 1 - cursor = span_end - - if message.get("step_loss_mask", 1) != 1: - continue - - if rendered_text[content_start : content_start + len(think_prefix)] == think_prefix: - mask_start = content_start + len(think_prefix) - else: - mask_start = content_start - - for pos in range(mask_start, span_end): - char_mask[pos] = 1 - - char_mask_prefix_sum = [0] - for value in char_mask: - char_mask_prefix_sum.append(char_mask_prefix_sum[-1] + value) - - loss_mask = [] - for start, end in offset_mapping: - if end <= start: - loss_mask.append(0) - else: - loss_mask.append(1 if char_mask_prefix_sum[end] - char_mask_prefix_sum[start] > 0 else 0) - return token_ids, loss_mask - def gen_multi_turn_loss_mask_gemma4( - self, messages: list[dict], tools: list[dict] = None + def get_loss_mask_from_rendered_qwen3_5( + self, + messages: list[dict], + rendered_text: str, + input_ids: list[int] | None = None, ) -> tuple[list[int], list[int]]: - """Mask assistant content plus ```` in Gemma4 chat templates.""" - rendered_text = self.tokenizer.apply_chat_template(messages, tokenize=False, tools=tools, return_dict=False) tokenized = self.tokenizer(rendered_text, add_special_tokens=False, return_offsets_mapping=True) token_ids = tokenized["input_ids"] offset_mapping = tokenized.get("offset_mapping") if offset_mapping is None: raise ValueError( - "Gemma4 loss mask generation requires a fast tokenizer with `return_offsets_mapping` support." + "Qwen3.5/Qwen3.6 loss mask generation requires a fast tokenizer with offset mapping support." ) - expected_token_ids = self.tokenizer.apply_chat_template( - messages, tokenize=True, tools=tools, return_dict=False - ) - if token_ids != expected_token_ids: + if input_ids is not None and token_ids != input_ids: + mismatch_idx = next( + (idx for idx, (actual, expected) in enumerate(zip(token_ids, input_ids)) if actual != expected), + min(len(token_ids), len(input_ids)), + ) raise ValueError( - "Gemma4 rendered text tokenization does not match " "`apply_chat_template(..., tokenize=True)` output." + "Rendered multimodal text tokenization does not match processor input_ids: " + f"{len(token_ids)=}, {len(input_ids)=}, first_mismatch_index={mismatch_idx}. " + "Please check processor visual-token expansion logic." ) - assistant_header = "<|turn>model\n" - think_open = "<|channel>thought\n" - think_close = "" - end_marker = "" + assistant_header = "<|im_start|>assistant\n" + think_prefix = "\n" + empty_think_block = "\n\n\n\n" + end_marker = "<|im_end|>" char_mask = [0] * len(rendered_text) cursor = 0 + def mark_span(start: int, end: int) -> None: + for pos in range(max(start, 0), min(end, len(char_mask))): + char_mask[pos] = 1 + for message in messages: if message["role"] != "assistant": continue header_pos = rendered_text.find(assistant_header, cursor) if header_pos < 0: - raise ValueError("Failed to locate assistant (model) turn in rendered Gemma4 chat template output.") + raise ValueError( + "Failed to locate assistant message in rendered Qwen3.5/Qwen3.6 chat template." + ) content_start = header_pos + len(assistant_header) end_pos = rendered_text.find(end_marker, content_start) if end_pos < 0: - raise ValueError("Failed to locate for assistant message in rendered Gemma4 text.") + raise ValueError("Failed to locate <|im_end|> for assistant message in rendered text.") span_end = end_pos + len(end_marker) if span_end < len(rendered_text) and rendered_text[span_end] == "\n": @@ -246,15 +227,19 @@ def gen_multi_turn_loss_mask_gemma4( if message.get("step_loss_mask", 1) != 1: continue - mask_start = content_start - if rendered_text[content_start : content_start + len(think_open)] == think_open: - close_pos = rendered_text.find(think_close, content_start) - if close_pos < 0: - raise ValueError("Found <|channel>thought open without matching close.") - mask_start = close_pos + len(think_close) + # In non-thinking mode, Qwen3.6's generation prompt pre-fills the + # empty think block. The model should learn only the answer/tool + # call after it, not generate a duplicate closing . + if rendered_text[content_start : content_start + len(empty_think_block)] == empty_think_block: + mask_start = content_start + len(empty_think_block) + # In thinking mode, only the opening \n is prompt scaffold; + # reasoning, , answer, and <|im_end|> are model targets. + elif rendered_text[content_start : content_start + len(think_prefix)] == think_prefix: + mask_start = content_start + len(think_prefix) + else: + mask_start = content_start - for pos in range(mask_start, span_end): - char_mask[pos] = 1 + mark_span(mask_start, span_end) char_mask_prefix_sum = [0] for value in char_mask: @@ -269,6 +254,93 @@ def gen_multi_turn_loss_mask_gemma4( return token_ids, loss_mask + @staticmethod + def _to_plain_list(value): + if value is None: + return [] + if hasattr(value, "detach"): + value = value.detach().cpu() + if hasattr(value, "tolist"): + value = value.tolist() + return value + + @staticmethod + def _grid_token_count(grid, merge_size: int) -> int: + prod = 1 + for dim in grid: + prod *= int(dim) + return prod // (int(merge_size) ** 2) + + def _expand_image_tokens_in_rendered_text( + self, + rendered_text: str, + image_grid_thw, + *, + image_token: str, + merge_size: int, + ) -> str: + grids = self._to_plain_list(image_grid_thw) + if not grids: + return rendered_text + + pieces = [] + cursor = 0 + for grid in grids: + image_pos = rendered_text.find(image_token, cursor) + if image_pos < 0: + raise ValueError( + "Failed to locate image token in rendered multimodal text while expanding visual tokens: " + f"{image_token=}, expanded_images={len(pieces)}." + ) + + num_image_tokens = self._grid_token_count(grid, merge_size) + pieces.append(rendered_text[cursor:image_pos]) + pieces.append(image_token * num_image_tokens) + cursor = image_pos + len(image_token) + + pieces.append(rendered_text[cursor:]) + return "".join(pieces) + + @staticmethod + def _message_content_to_text(content) -> str: + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, list): + text_parts = [] + for item in content: + if isinstance(item, str): + text_parts.append(item) + elif isinstance(item, dict): + if "text" in item: + text_parts.append(item["text"]) + elif item.get("type") == "text": + text_parts.append(item.get("text", "")) + return "".join(text_parts) + return str(content) + + def _get_assistant_content_parts(self, message: dict) -> list[str]: + content = self._message_content_to_text(message.get("content")).strip() + reasoning_content = message.get("reasoning_content") + parts = [] + + if isinstance(reasoning_content, str): + reasoning_content = reasoning_content.strip() + if reasoning_content: + parts.append(reasoning_content) + elif "" in content: + reasoning_content = content.split("", 1)[0].rstrip("\n").split("")[-1].lstrip("\n").strip() + if reasoning_content: + parts.append(reasoning_content) + content = content.split("", 1)[-1].lstrip("\n") + + content = content.strip() + if content: + parts.append(content) + + return parts + def gen_multi_turn_loss_mask_distill_qwen( self, messages: list[dict], tools: list[dict] = None ) -> tuple[list[int], list[int]]: @@ -305,31 +377,32 @@ def get_loss_mask(self, messages: list[dict], tools: list[dict] = None) -> tuple raise ValueError(f"Unsupported tokenizer type: {self.tokenizer_type}") def get_loss_mask_with_multimodal_alignment( - self, messages: list[dict], input_ids: list[int], tools: list[dict] = None + self, + messages: list[dict], + input_ids: list[int], + tools: list[dict] = None, + rendered_text: str | None = None, + image_grid_thw=None, + image_token: str = "<|image_pad|>", + image_merge_size: int = 2, ) -> tuple[list[int], list[int]]: - text = [] - for msg in messages: - if isinstance(msg.get("content"), list): - text_parts = [] - for item in msg["content"]: - if isinstance(item, dict) and item.get("type") == "text": - text_parts.append(item.get("text", "")) - elif isinstance(item, str): - text_parts.append(item) - text.append({"role": msg["role"], "content": " ".join(text_parts)}) - else: - text.append(msg) - - _, loss_mask_text = self.get_loss_mask(text, tools=tools) + if self.tokenizer_type == "qwen3_5" and rendered_text is not None: + expanded_rendered_text = self._expand_image_tokens_in_rendered_text( + rendered_text, + image_grid_thw, + image_token=image_token, + merge_size=image_merge_size, + ) + return self.get_loss_mask_from_rendered_qwen3_5( + messages, + expanded_rendered_text, + input_ids=input_ids, + ) - diff = len(input_ids) - len(loss_mask_text) - assert diff >= 0, ( - f"input_ids (length={len(input_ids)}) is shorter than text loss_mask (length={len(loss_mask_text)}) " - f"Please check if processor and tokenizer tokenization are consistent." + raise ValueError( + "Multimodal loss-mask alignment requires the actual rendered multimodal text for Qwen3.5/Qwen3.6. " + "Pass rendered_text plus processor image_grid_thw so visual token positions can be masked directly." ) - loss_mask = [0] * diff + loss_mask_text - - return input_ids, loss_mask def get_text_from_loss_mask(self, token_ids: list[int], loss_masks: list[int]) -> list[str]: selected_texts = [] diff --git a/slime/utils/ppo_utils.py b/slime/utils/ppo_utils.py index cfce0324c4..ef009e4b77 100644 --- a/slime/utils/ppo_utils.py +++ b/slime/utils/ppo_utils.py @@ -858,3 +858,44 @@ def calculate_log_probs_and_entropy( entropy = logits.new_zeros((0,)) return log_prob, entropy + + +def calculate_log_probs_and_entropy_fused(hidden, weight, tokens, tp_group, with_entropy: bool = False, chunk_size: int = -1): + """Compute LM-head logits and CE in chunks to avoid materializing full [T, V].""" + import torch.nn.functional as F + from torch.utils.checkpoint import checkpoint as _ckpt + + hidden = hidden.contiguous() + entropy = None + T = hidden.size(0) + if T == 0: + log_prob = hidden.new_zeros((0, 1), dtype=torch.float32) + if with_entropy: + entropy = hidden.new_zeros((0,), dtype=torch.float32) + return log_prob, entropy + + chunk = chunk_size if chunk_size and chunk_size > 0 else T + + def _chunk_fn(hidden_chunk, token_chunk, lm_head_weight): + logits_chunk = F.linear(hidden_chunk.to(lm_head_weight.dtype), lm_head_weight).float() + return _calculate_log_probs_and_entropy_chunk( + logits_chunk, + token_chunk, + tp_group, + with_entropy=with_entropy, + log_prob_keep_mask=None, + ) + + log_probs, entropys = [], [] + for start in range(0, T, chunk): + hidden_chunk = hidden[start : start + chunk] + token_chunk = tokens[start : start + chunk] + log_prob_chunk, entropy_chunk = _ckpt(_chunk_fn, hidden_chunk, token_chunk, weight, use_reentrant=False) + log_probs.append(log_prob_chunk) + if with_entropy: + entropys.append(entropy_chunk) + + log_prob = torch.cat(log_probs, dim=0) + if with_entropy: + entropy = torch.cat(entropys, dim=0) + return log_prob, entropy diff --git a/slime/utils/processing_utils.py b/slime/utils/processing_utils.py index fb652e73d0..5d81f8f31e 100644 --- a/slime/utils/processing_utils.py +++ b/slime/utils/processing_utils.py @@ -3,16 +3,28 @@ import json import logging from pathlib import Path +import time +from urllib.parse import urlparse -from PIL import Image +from PIL import Image, ImageFile +import requests from transformers import AutoProcessor, AutoTokenizer, PreTrainedTokenizerBase, ProcessorMixin logger = logging.getLogger(__name__) +# Some generated PPT screenshots are mildly truncated but still decodable. +# Enable this before qwen_vl_utils opens image paths and returns lazy PIL objects. +ImageFile.LOAD_TRUNCATED_IMAGES = True + # Default image patch size for vision-language models # Note: Qwen3-VL uses 16, Qwen2.5-VL uses 14 # Reference: https://github.com/QwenLM/Qwen3-VL/blob/main/qwen-vl-utils/README.md DEFAULT_PATCH_SIZE = 14 +IMAGE_DOWNLOAD_TIMEOUT = (5, 30) +IMAGE_DOWNLOAD_RETRIES = 3 +IMAGE_DOWNLOAD_HEADERS = { + "User-Agent": "slime-vlm-data-loader/1.0", +} def load_tokenizer(name_or_path: str, **kwargs): @@ -97,6 +109,75 @@ def load_processor(name_or_path: str, **kwargs): return proc +def _is_url(value: str) -> bool: + parsed = urlparse(value) + return parsed.scheme in {"http", "https"} and bool(parsed.netloc) + + +def _load_image_from_bytes(raw: bytes, source: str) -> Image.Image: + try: + image = Image.open(io.BytesIO(raw)) + image.load() + return image + except Exception as e: + raise RuntimeError(f"Failed to decode image from {source}") from e + + +def _download_image(url: str) -> Image.Image: + urls = [url] + if url.startswith("https://"): + urls.append("http://" + url[len("https://") :]) + + last_error = None + for candidate in urls: + for attempt in range(1, IMAGE_DOWNLOAD_RETRIES + 1): + try: + response = requests.get( + candidate, + headers=IMAGE_DOWNLOAD_HEADERS, + timeout=IMAGE_DOWNLOAD_TIMEOUT, + ) + response.raise_for_status() + return _load_image_from_bytes(response.content, candidate) + except Exception as e: + last_error = e + if attempt < IMAGE_DOWNLOAD_RETRIES: + time.sleep(0.5 * attempt) + continue + logger.warning( + "Failed to load image URL %s after %d attempts: %s", + candidate, + IMAGE_DOWNLOAD_RETRIES, + e, + ) + + raise RuntimeError(f"Failed to load image URL {url}") from last_error + + +def _load_image_from_file(path: str) -> Image.Image: + try: + image = Image.open(path) + image.load() + return image + except Exception as e: + raise RuntimeError(f"Failed to decode image from file {path}") from e + + +def _load_image_from_string(image_data: str) -> Image.Image: + if image_data.startswith("data:"): + _, encoded = image_data.split(",", 1) + return _load_image_from_bytes(base64.b64decode(encoded), "data URI") + + if _is_url(image_data): + return _download_image(image_data) + + try: + raw = base64.b64decode(image_data) + return _load_image_from_bytes(raw, "base64 string") + except Exception: + return _load_image_from_file(image_data) + + def _extract_images_from_messages(messages): """Extract PIL images from chat messages containing multimodal content. @@ -117,16 +198,7 @@ def _extract_images_from_messages(messages): if isinstance(image_data, Image.Image): images.append(image_data) elif isinstance(image_data, str): - if image_data.startswith("data:"): - _, encoded = image_data.split(",", 1) - images.append(Image.open(io.BytesIO(base64.b64decode(encoded)))) - else: - try: - raw = base64.b64decode(image_data) - images.append(Image.open(io.BytesIO(raw))) - except Exception: - # Not base64 — try as file path - images.append(Image.open(image_data)) + images.append(_load_image_from_string(image_data)) return images diff --git a/slime_plugins/models/hf_attention.py b/slime_plugins/models/hf_attention.py index d5cf619724..833a5939ef 100644 --- a/slime_plugins/models/hf_attention.py +++ b/slime_plugins/models/hf_attention.py @@ -1,5 +1,6 @@ import json import os +import time from abc import ABC, abstractmethod import torch @@ -10,6 +11,55 @@ from megatron.core.transformer.module import MegatronModule +def _env_flag(name: str) -> bool: + value = os.environ.get(name) + if value is None: + value = os.environ.get(name.lower(), "0") + return value.lower() in {"1", "true", "yes", "on"} + + +def _env_int(name: str, default: int) -> int: + try: + value = os.environ.get(name) + if value is None: + value = os.environ.get(name.lower(), str(default)) + return int(value) + except ValueError: + return default + + +def _global_rank() -> int: + if dist.is_available() and dist.is_initialized(): + return dist.get_rank() + return int(os.environ.get("RANK", "0")) + + +class _CudaStepTimer: + def __init__(self, enabled: bool): + self.enabled = enabled and torch.cuda.is_available() + self.events: list[tuple[str, torch.cuda.Event, torch.cuda.Event]] = [] + self.start_wall = time.perf_counter() + + def run(self, name, fn): + if not self.enabled: + return fn() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + result = fn() + end.record() + self.events.append((name, start, end)) + return result + + def summary(self) -> dict[str, float]: + if not self.enabled: + return {"wall_ms": (time.perf_counter() - self.start_wall) * 1000.0} + torch.cuda.synchronize() + summary = {name: start.elapsed_time(end) for name, start, end in self.events} + summary["total_ms"] = sum(summary.values()) + return summary + + def _load_hf_config(checkpoint_path): """Load HF config with fallback for unsupported model types.""" try: @@ -84,6 +134,29 @@ def __init__( self.hf_config = _load_hf_config(args.hf_checkpoint) # hardcode to fa2 at the moment. self.hf_config._attn_implementation = "flash_attention_2" + if not getattr(HuggingfaceAttention, "_slime_hf_attention_timing_logged", False): + HuggingfaceAttention._slime_hf_attention_timing_logged = True + print( + "SLIME HF attention timing config: " + f"enabled={_env_flag('SLIME_HF_ATTENTION_TIMING')} " + f"rank={_env_int('SLIME_HF_ATTENTION_TIMING_RANK', 0)} " + f"interval={_env_int('SLIME_HF_ATTENTION_TIMING_INTERVAL', 32)} " + f"limit={_env_int('SLIME_HF_ATTENTION_TIMING_LIMIT', 64)} " + f"file={__file__}", + flush=True, + ) + + def _timing_should_sample(self) -> bool: + if not _env_flag("SLIME_HF_ATTENTION_TIMING"): + return False + rank_filter = _env_int("SLIME_HF_ATTENTION_TIMING_RANK", 0) + if rank_filter >= 0 and _global_rank() != rank_filter: + return False + count = getattr(self, "_slime_hf_attention_timing_count", 0) + 1 + self._slime_hf_attention_timing_count = count + limit = _env_int("SLIME_HF_ATTENTION_TIMING_LIMIT", 64) + interval = max(1, _env_int("SLIME_HF_ATTENTION_TIMING_INTERVAL", 32)) + return count <= limit and (count == 1 or count % interval == 0) def forward( self, @@ -101,69 +174,105 @@ def forward( *, inference_params: BaseInferenceContext | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - assert packed_seq_params is not None - cu_seqlens = packed_seq_params.cu_seqlens_q + cu_seqlens = packed_seq_params.cu_seqlens_q if packed_seq_params is not None else None + should_time = self._timing_should_sample() + timer = _CudaStepTimer(should_time) + timing_meta = { + "layer": self.layer_number, + "input_shape": tuple(hidden_states.shape), + "cp": mpu.get_context_parallel_world_size(), + "sp": bool(self.args.sequence_parallel), + "packed": cu_seqlens is not None, + } if self.args.sequence_parallel: # tensor_parallel_output_grad=False: the linear attention after this # gather is NOT TP-sharded (duplicated on all ranks), so the backward # should split (not reduce-scatter) to avoid inflating gradients by TP. - hidden_states = tensor_parallel.gather_from_sequence_parallel_region( - hidden_states, - tensor_parallel_output_grad=False, - group=mpu.get_tensor_model_parallel_group(), + hidden_states = timer.run( + "sp_gather", + lambda: tensor_parallel.gather_from_sequence_parallel_region( + hidden_states, + tensor_parallel_output_grad=False, + group=mpu.get_tensor_model_parallel_group(), + ), ) if mpu.get_context_parallel_world_size() > 1: + assert cu_seqlens is not None, "Context parallelism for HF attention requires packed_seq_params." cp_size = mpu.get_context_parallel_world_size() # Use custom all-gather whose backward returns local gradient # instead of reduce-scatter, since the computation is duplicated. - hidden_states_list = _AllGatherForDuplicatedComputation.apply( - hidden_states, - mpu.get_context_parallel_group(), + hidden_states_list = timer.run( + "cp_all_gather", + lambda: _AllGatherForDuplicatedComputation.apply( + hidden_states, + mpu.get_context_parallel_group(), + ), ) # TODO: preprocess this for each batch to prevent tolist in the training step - whole_hidden_states_list = [] - - local_cu_seqlens = cu_seqlens // cp_size - for i in range(len(cu_seqlens) - 1): - seqlen = cu_seqlens[i + 1] - cu_seqlens[i] - chunk_size = seqlen // 2 // cp_size - whole_hidden_states_list.extend( - [ - hidden_states_list[cp_rank][local_cu_seqlens[i] : local_cu_seqlens[i] + chunk_size] - for cp_rank in range(cp_size) - ] - + [ - hidden_states_list[cp_rank][local_cu_seqlens[i] + chunk_size : local_cu_seqlens[i + 1]] - for cp_rank in range(cp_size) - ][::-1], - ) - hidden_states = torch.cat(whole_hidden_states_list, dim=0) - - hidden_states = hidden_states.permute(1, 0, 2) # [bsz, seq_len, hidden_dim] - - output = self.hf_forward(hidden_states, packed_seq_params) + def rebuild_cp_hidden_states(): + whole_hidden_states_list = [] + local_cu_seqlens = cu_seqlens // cp_size + for i in range(len(cu_seqlens) - 1): + seqlen = cu_seqlens[i + 1] - cu_seqlens[i] + chunk_size = seqlen // 2 // cp_size + whole_hidden_states_list.extend( + [ + hidden_states_list[cp_rank][local_cu_seqlens[i] : local_cu_seqlens[i] + chunk_size] + for cp_rank in range(cp_size) + ] + + [ + hidden_states_list[cp_rank][local_cu_seqlens[i] + chunk_size : local_cu_seqlens[i + 1]] + for cp_rank in range(cp_size) + ][::-1], + ) + return torch.cat(whole_hidden_states_list, dim=0) + + hidden_states = timer.run("cp_rebuild", rebuild_cp_hidden_states) + + hidden_states = timer.run("permute_in", lambda: hidden_states.permute(1, 0, 2)) # [bsz, seq_len, hidden_dim] + + output = timer.run("hf_forward", lambda: self.hf_forward(hidden_states, packed_seq_params)) bias = None - output = output.permute(1, 0, 2) # [seq_len, bsz, hidden_dim] + output = timer.run("permute_out", lambda: output.permute(1, 0, 2)) # [seq_len, bsz, hidden_dim] if mpu.get_context_parallel_world_size() > 1: cp_rank = mpu.get_context_parallel_rank() - output_list = [] - for i in range(len(cu_seqlens) - 1): - seqlen = cu_seqlens[i + 1] - cu_seqlens[i] - chunk_size = seqlen // 2 // cp_size - seq = output[cu_seqlens[i] : cu_seqlens[i + 1]] - chunks = torch.chunk(seq, 2 * cp_size, dim=0) - output_list.append(chunks[cp_rank]) - output_list.append(chunks[2 * cp_size - 1 - cp_rank]) - output = torch.cat(output_list, dim=0) + + def scatter_cp_output(): + output_list = [] + for i in range(len(cu_seqlens) - 1): + seqlen = cu_seqlens[i + 1] - cu_seqlens[i] + chunk_size = seqlen // 2 // cp_size + seq = output[cu_seqlens[i] : cu_seqlens[i + 1]] + chunks = torch.chunk(seq, 2 * cp_size, dim=0) + output_list.append(chunks[cp_rank]) + output_list.append(chunks[2 * cp_size - 1 - cp_rank]) + return torch.cat(output_list, dim=0) + + output = timer.run("cp_scatter", scatter_cp_output) if self.args.sequence_parallel: - output = tensor_parallel.scatter_to_sequence_parallel_region( - output, group=mpu.get_tensor_model_parallel_group() + output = timer.run( + "sp_scatter", + lambda: tensor_parallel.scatter_to_sequence_parallel_region( + output, group=mpu.get_tensor_model_parallel_group() + ), + ) + + if should_time: + if cu_seqlens is not None: + timing_meta["num_sequences"] = int(cu_seqlens.numel() - 1) + timing_meta["max_seqlen"] = int((cu_seqlens[1:] - cu_seqlens[:-1]).max().detach().item()) + timing_meta["output_shape"] = tuple(output.shape) + print( + "SLIME HF attention timing: " + f"meta={timing_meta} times_ms=" + f"{ {key: round(value, 3) for key, value in timer.summary().items()} }", + flush=True, ) return output, bias diff --git a/slime_plugins/models/qwen3_5.py b/slime_plugins/models/qwen3_5.py index 294c9d97a9..de165af477 100644 --- a/slime_plugins/models/qwen3_5.py +++ b/slime_plugins/models/qwen3_5.py @@ -1,4 +1,5 @@ import copy +import logging import torch import torch.nn as nn @@ -18,6 +19,10 @@ from .qwen_gdn_backend import get_chunk_gated_delta_rule +logger = logging.getLogger(__name__) +_LOGGED_GDN_BACKENDS = set() + + def _get_text_config(hf_config): """Extract text config from a VLM config if needed.""" if hasattr(hf_config, "text_config"): @@ -37,6 +42,13 @@ def __init__(self, config, layer_idx: int, args=None): super().__init__() self.gdn_backend = getattr(args, "qwen_gdn_backend", "fla") self.chunk_gated_delta_rule = get_chunk_gated_delta_rule(self.gdn_backend) + if self.gdn_backend not in _LOGGED_GDN_BACKENDS: + logger.info( + "Qwen3.5 GDN backend selected: %s%s", + self.gdn_backend, + " (auto_cp=True)" if self.gdn_backend == "flashqla" else "", + ) + _LOGGED_GDN_BACKENDS.add(self.gdn_backend) self.hidden_size = config.hidden_size self.num_v_heads = config.linear_num_value_heads self.num_k_heads = config.linear_num_key_heads @@ -127,6 +139,10 @@ def forward( g = g.contiguous() beta = beta.contiguous() + kernel_kwargs = {} + if self.gdn_backend == "flashqla": + kernel_kwargs["auto_cp"] = True + core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule( query, key, @@ -137,6 +153,7 @@ def forward( output_final_state=False, use_qk_l2norm_in_kernel=True, cu_seqlens=cu_seqlens, + **kernel_kwargs, ) z_shape_og = z.shape @@ -185,9 +202,10 @@ def __init__( def hf_forward(self, hidden_states, packed_seq_params): hidden_states = self.input_layernorm(hidden_states) + cu_seqlens = packed_seq_params.cu_seqlens_q if packed_seq_params is not None else None hidden_states = self.linear_attn( hidden_states=hidden_states, - cu_seqlens=packed_seq_params.cu_seqlens_q, + cu_seqlens=cu_seqlens, ) return hidden_states diff --git a/slime_plugins/models/qwen3_next.py b/slime_plugins/models/qwen3_next.py index 683cb28063..b40a38a87e 100644 --- a/slime_plugins/models/qwen3_next.py +++ b/slime_plugins/models/qwen3_next.py @@ -1,4 +1,5 @@ import copy +import logging import torch import torch.nn as nn @@ -21,6 +22,10 @@ from .qwen_gdn_backend import get_chunk_gated_delta_rule +logger = logging.getLogger(__name__) +_LOGGED_GDN_BACKENDS = set() + + # adapt from https://github.com/huggingface/transformers/blob/38a08b6e8ae35857109cedad75377997fecbf9d0/src/transformers/models/qwen3_next/modeling_qwen3_next.py#L564 class Qwen3NextGatedDeltaNet(nn.Module): """ @@ -31,6 +36,13 @@ def __init__(self, config, layer_idx: int, args=None): super().__init__() self.gdn_backend = getattr(args, "qwen_gdn_backend", "fla") self.chunk_gated_delta_rule = get_chunk_gated_delta_rule(self.gdn_backend) + if self.gdn_backend not in _LOGGED_GDN_BACKENDS: + logger.info( + "Qwen3Next GDN backend selected: %s%s", + self.gdn_backend, + " (auto_cp=True)" if self.gdn_backend == "flashqla" else "", + ) + _LOGGED_GDN_BACKENDS.add(self.gdn_backend) self.hidden_size = config.hidden_size self.num_v_heads = config.linear_num_value_heads self.num_k_heads = config.linear_num_key_heads @@ -149,6 +161,10 @@ def forward( g = g.contiguous() beta = beta.contiguous() + kernel_kwargs = {} + if self.gdn_backend == "flashqla": + kernel_kwargs["auto_cp"] = True + core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule( query, key, @@ -159,6 +175,7 @@ def forward( output_final_state=False, use_qk_l2norm_in_kernel=True, cu_seqlens=cu_seqlens, + **kernel_kwargs, ) z_shape_og = z.shape diff --git a/tests/test_dp_schedule.py b/tests/test_dp_schedule.py index 33a8be53b5..a6ce0fd294 100644 --- a/tests/test_dp_schedule.py +++ b/tests/test_dp_schedule.py @@ -19,23 +19,19 @@ def make_args( micro_batch_size=1, use_dynamic_batch_size=False, max_tokens_per_gpu=None, + packing_safety_margin=1.0, + global_batch_tokens=None, balance_data=False, - balance_by_flops=False, + multimodal_aware_packing="off", ): return SimpleNamespace( micro_batch_size=micro_batch_size, use_dynamic_batch_size=use_dynamic_batch_size, max_tokens_per_gpu=max_tokens_per_gpu, + packing_safety_margin=packing_safety_margin, + global_batch_tokens=global_batch_tokens, balance_data=balance_data, - balance_by_flops=balance_by_flops, - hidden_size=16, - num_attention_heads=2, - num_query_groups=2, - vocab_size=32, - ffn_hidden_size=64, - num_experts=None, - num_layers=2, - kv_channels=8, + multimodal_aware_packing=multimodal_aware_packing, ) @@ -224,6 +220,152 @@ def test_dynamic_oversized_sample_lands_alone(): assert found +@pytest.mark.unit +def test_multimodal_aware_packing_separates_text_and_mm_samples(): + """Opt-in multimodal-aware packing avoids mixing true image samples with text-only fastpath samples.""" + total_lengths = [70, 20, 20, 20, 20, 20, 20, 20] + rollout_indices = list(range(8)) + sample_is_multimodal = [True, False, False, False, False, False, False, False] + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=100, multimodal_aware_packing="separate") + tp = make_tp(dp_size=2) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, + tp, + total_lengths, + global_batch_size=8, + rollout_indices=rollout_indices, + sample_is_multimodal=sample_is_multimodal, + ) + + assert gbs_per_step == [8] + assert_invariants( + partitions, + mbi, + nmb, + dp_size=2, + expected_global_sample_indices=range(8), + total_lengths=total_lengths, + max_per_bin=100, + ) + for r in range(2): + for mbs in mbi[r]: + flags = {sample_is_multimodal[partitions[r][i]] for i in mbs} + assert len(flags) == 1 + + +@pytest.mark.unit +def test_multimodal_aware_packing_uses_padded_unsplit_cost_for_mm_bins(): + """Two image samples whose raw sum fits should split when padded unsplit cost would exceed the cap.""" + total_lengths = [80, 20, 10, 10, 10, 10, 10, 10] + rollout_indices = list(range(8)) + sample_is_multimodal = [True, True, False, False, False, False, False, False] + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=100, multimodal_aware_packing="separate") + tp = make_tp(dp_size=2) + + partitions, mbi, nmb, _ = build_dp_schedule( + args, + tp, + total_lengths, + global_batch_size=8, + rollout_indices=rollout_indices, + sample_is_multimodal=sample_is_multimodal, + ) + + mm_pair_seen = False + for r in range(2): + for mbs in mbi[r]: + globals_ = [partitions[r][i] for i in mbs] + if 0 in globals_ and 1 in globals_: + mm_pair_seen = True + assert not mm_pair_seen + assert_invariants( + partitions, + mbi, + nmb, + dp_size=2, + expected_global_sample_indices=range(8), + total_lengths=total_lengths, + max_per_bin=100, + ) + + +@pytest.mark.unit +def test_multimodal_aware_packing_separate_raw_keeps_raw_sum_cost_for_mm_bins(): + """separate_raw keeps text/mm separation without adding padded unsplit cost.""" + total_lengths = [80, 20, 10, 10, 10, 10, 10, 10] + rollout_indices = list(range(8)) + sample_is_multimodal = [True, True, False, False, False, False, False, False] + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=100, multimodal_aware_packing="separate_raw") + tp = make_tp(dp_size=2) + + partitions, mbi, nmb, _ = build_dp_schedule( + args, + tp, + total_lengths, + global_batch_size=8, + rollout_indices=rollout_indices, + sample_is_multimodal=sample_is_multimodal, + ) + + mm_pair_seen = False + for r in range(2): + for mbs in mbi[r]: + globals_ = [partitions[r][i] for i in mbs] + flags = {sample_is_multimodal[g] for g in globals_} + assert len(flags) == 1 + if 0 in globals_ and 1 in globals_: + mm_pair_seen = True + assert mm_pair_seen + assert_invariants( + partitions, + mbi, + nmb, + dp_size=2, + expected_global_sample_indices=range(8), + total_lengths=total_lengths, + max_per_bin=100, + ) + + +@pytest.mark.unit +def test_multimodal_aware_packing_separate_raw_uses_raw_cost_for_mm_bins(): + """separate_raw avoids text/mm mixing but does not split mm samples by padded unsplit cost.""" + total_lengths = [80, 20, 10, 10, 10, 10, 10, 10] + rollout_indices = list(range(8)) + sample_is_multimodal = [True, True, False, False, False, False, False, False] + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=100, multimodal_aware_packing="separate_raw") + tp = make_tp(dp_size=2) + + partitions, mbi, nmb, _ = build_dp_schedule( + args, + tp, + total_lengths, + global_batch_size=8, + rollout_indices=rollout_indices, + sample_is_multimodal=sample_is_multimodal, + ) + + mm_pair_seen = False + for r in range(2): + for mbs in mbi[r]: + globals_ = [partitions[r][i] for i in mbs] + flags = {sample_is_multimodal[g] for g in globals_} + assert len(flags) == 1 + if 0 in globals_ and 1 in globals_: + mm_pair_seen = True + assert mm_pair_seen + assert_invariants( + partitions, + mbi, + nmb, + dp_size=2, + expected_global_sample_indices=range(8), + total_lengths=total_lengths, + max_per_bin=100, + ) + + @pytest.mark.unit def test_dynamic_with_vpp_rounds_to_mb_group(): """num_microbatches per rank should be a multiple of mb_group when vpp_size > 1.""" @@ -313,6 +455,53 @@ def test_trims_trailing_rollouts_that_dont_fill_a_step(): ) +@pytest.mark.unit +def test_token_based_global_batch_uses_token_budget(): + total_lengths = [8, 8, 8, 8, 2, 2, 2, 2] + rollout_indices = list(range(8)) + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=12, global_batch_tokens=20) + tp = make_tp(dp_size=2) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=8, rollout_indices=rollout_indices + ) + + assert gbs_per_step == [2, 4, 2] + assert_invariants( + partitions, + mbi, + nmb, + dp_size=2, + expected_global_sample_indices=range(8), + total_lengths=total_lengths, + max_per_bin=12, + ) + + +@pytest.mark.unit +def test_packing_safety_margin_plans_more_microbatches(): + total_lengths = [6, 6, 6, 6] + rollout_indices = list(range(4)) + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=12, packing_safety_margin=0.5) + tp = make_tp(dp_size=1) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=4, rollout_indices=rollout_indices + ) + + assert gbs_per_step == [4] + assert nmb == [4] + assert_invariants( + partitions, + mbi, + nmb, + dp_size=1, + expected_global_sample_indices=range(4), + total_lengths=total_lengths, + max_per_bin=12, + ) + + @pytest.mark.unit def test_rejects_when_fewer_rollouts_than_gbs(): """gbs=4 with only 3 distinct rollouts → cannot form one step.""" diff --git a/tests/utils/test_loss_mask_type_qwen35.py b/tests/utils/test_loss_mask_type_qwen35.py index 69ffd1513f..e93bdeddfe 100644 --- a/tests/utils/test_loss_mask_type_qwen35.py +++ b/tests/utils/test_loss_mask_type_qwen35.py @@ -7,7 +7,8 @@ class FakeQwen35Tokenizer: The critical behavior we need for these tests is: 1. `add_generation_prompt=True` appends `<|im_start|>assistant\n\n` 2. Only assistant turns after the last non-tool user query are wrapped in - `...` + `...`, unless `preserve_thinking=True` is passed through + `apply_chat_template`. """ assistant_generation_prompt = "<|im_start|>assistant\n\n" @@ -31,16 +32,24 @@ def apply_chat_template( add_special_tokens=False, **kwargs, ): - rendered = self.render(messages, add_generation_prompt=add_generation_prompt) + rendered = self.render( + messages, + add_generation_prompt=add_generation_prompt, + preserve_thinking=kwargs.get("preserve_thinking", False), + ) if tokenize: return [ord(ch) for ch in rendered] return rendered - def render(self, messages, add_generation_prompt=False): - rendered, _ = self.render_with_expected_mask(messages, add_generation_prompt=add_generation_prompt) + def render(self, messages, add_generation_prompt=False, preserve_thinking=False): + rendered, _ = self.render_with_expected_mask( + messages, + add_generation_prompt=add_generation_prompt, + preserve_thinking=preserve_thinking, + ) return rendered - def render_with_expected_mask(self, messages, add_generation_prompt=False): + def render_with_expected_mask(self, messages, add_generation_prompt=False, preserve_thinking=False): if not messages: raise ValueError("No messages provided.") @@ -56,13 +65,13 @@ def render_with_expected_mask(self, messages, add_generation_prompt=False): if role == "system": if index != 0: raise ValueError("System message must be at the beginning.") - piece = f"<|im_start|>system\n{message['content']}<|im_end|>\n" + piece = f"<|im_start|>system\n{self._render_content(message['content'])}<|im_end|>\n" pieces.append(piece) mask.extend([0] * len(piece)) continue if role == "user": - piece = f"<|im_start|>user\n{message['content']}<|im_end|>\n" + piece = f"<|im_start|>user\n{self._render_content(message['content'])}<|im_end|>\n" pieces.append(piece) mask.extend([0] * len(piece)) continue @@ -71,7 +80,7 @@ def render_with_expected_mask(self, messages, add_generation_prompt=False): piece = "" if index > 0 and messages[index - 1]["role"] != "tool": piece += "<|im_start|>user" - piece += f"\n\n{message['content']}\n" + piece += f"\n\n{self._render_content(message['content'])}\n" if index == len(messages) - 1 or messages[index + 1]["role"] != "tool": piece += "<|im_end|>\n" pieces.append(piece) @@ -82,19 +91,28 @@ def render_with_expected_mask(self, messages, add_generation_prompt=False): raise NotImplementedError(f"Unsupported role in test tokenizer: {role}") reasoning, answer = self._split_assistant_content(message["content"]) - if index > last_query_index: - prefix = "<|im_start|>assistant\n\n" - target = f"{reasoning}\n\n\n{answer}" + assistant_header = "<|im_start|>assistant\n" + if preserve_thinking or index > last_query_index: + body = f"\n{reasoning}\n\n\n{answer}" else: - prefix = "<|im_start|>assistant\n" - target = answer - - target += self._render_tool_calls(answer, message.get("tool_calls")) - target += "<|im_end|>\n" + body = answer + + body += self._render_tool_calls(answer, message.get("tool_calls")) + piece = assistant_header + body + "<|im_end|>\n" + + content_start = len(assistant_header) + empty_think_block = "\n\n\n\n" + think_prefix = "\n" + if piece[content_start : content_start + len(empty_think_block)] == empty_think_block: + mask_start = content_start + len(empty_think_block) + elif piece[content_start : content_start + len(think_prefix)] == think_prefix: + mask_start = content_start + len(think_prefix) + else: + mask_start = content_start - pieces.append(prefix + target) - mask.extend([0] * len(prefix)) - mask.extend([1] * len(target)) + pieces.append(piece) + mask.extend([0] * mask_start) + mask.extend([1] * (len(piece) - mask_start)) if add_generation_prompt: pieces.append(self.assistant_generation_prompt) @@ -102,6 +120,28 @@ def render_with_expected_mask(self, messages, add_generation_prompt=False): return "".join(pieces), mask + @staticmethod + def _render_content(content): + if isinstance(content, str): + return content + if isinstance(content, list): + pieces = [] + for item in content: + if isinstance(item, str): + pieces.append(item) + elif isinstance(item, dict) and item.get("type") == "image": + pieces.append("<|vision_start|><|image_pad|><|vision_end|>") + elif isinstance(item, dict) and "image" in item: + pieces.append("<|vision_start|><|image_pad|><|vision_end|>") + elif isinstance(item, dict) and item.get("type") == "text": + pieces.append(item.get("text", "")) + elif isinstance(item, dict) and "text" in item: + pieces.append(item["text"]) + else: + raise NotImplementedError(f"Unsupported content item in test tokenizer: {item}") + return "".join(pieces) + return "" if content is None else str(content) + @staticmethod def _split_assistant_content(content): if "" not in content: @@ -156,11 +196,73 @@ def test_qwen3_and_qwen3_5_match_on_single_turn_qwen35_data(): expected_text, expected_mask = tokenizer.render_with_expected_mask(messages) expected_token_ids = tokenizer(expected_text, add_special_tokens=False)["input_ids"] - for loss_mask_type in ("qwen3", "qwen3_5"): - generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type=loss_mask_type) - token_ids, loss_mask = generator.get_loss_mask(messages) - assert token_ids == expected_token_ids - assert loss_mask == expected_mask + qwen3_generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type="qwen3") + qwen35_generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type="qwen3_5") + + qwen3_token_ids, qwen3_loss_mask = qwen3_generator.get_loss_mask(messages) + qwen35_token_ids, qwen35_loss_mask = qwen35_generator.get_loss_mask(messages) + + assert qwen3_token_ids == expected_token_ids + assert qwen35_token_ids == expected_token_ids + assert qwen3_loss_mask == expected_mask + assert qwen35_loss_mask == expected_mask + assert qwen3_generator.get_text_from_loss_mask(qwen3_token_ids, qwen3_loss_mask) == [ + "REASONING\n\n\nANSWER<|im_end|>\n", + ] + assert qwen35_generator.get_text_from_loss_mask(qwen35_token_ids, qwen35_loss_mask) == [ + "REASONING\n\n\nANSWER<|im_end|>\n", + ] + + +def test_qwen3_5_empty_think_block_is_prompt_scaffold(): + tokenizer = FakeQwen35Tokenizer() + messages = [ + {"role": "user", "content": "USER"}, + {"role": "assistant", "content": "ANSWER"}, + ] + + expected_text, expected_mask = tokenizer.render_with_expected_mask(messages) + expected_token_ids = tokenizer(expected_text, add_special_tokens=False)["input_ids"] + + generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type="qwen3_5") + token_ids, loss_mask = generator.get_loss_mask(messages) + + assert token_ids == expected_token_ids + assert loss_mask == expected_mask + assert generator.get_text_from_loss_mask(token_ids, loss_mask) == [ + "ANSWER<|im_end|>\n", + ] + assert "" not in generator.get_text_from_loss_mask(token_ids, loss_mask)[0] + + +def test_qwen3_5_passes_chat_template_kwargs_to_preserve_historical_thinking(): + tokenizer = FakeQwen35Tokenizer() + messages = [ + {"role": "user", "content": "USER_1"}, + {"role": "assistant", "content": "REASONING_1\nANSWER_1"}, + {"role": "user", "content": "USER_2"}, + {"role": "assistant", "content": "REASONING_2\nANSWER_2"}, + ] + + default_text, default_mask = tokenizer.render_with_expected_mask(messages) + preserved_text, preserved_mask = tokenizer.render_with_expected_mask(messages, preserve_thinking=True) + assert "REASONING_1" not in default_text + assert "REASONING_1" in preserved_text + + generator = MultiTurnLossMaskGenerator( + tokenizer, + tokenizer_type="qwen3_5", + chat_template_kwargs={"preserve_thinking": True}, + ) + token_ids, loss_mask = generator.get_loss_mask(messages) + + assert token_ids == tokenizer(preserved_text, add_special_tokens=False)["input_ids"] + assert loss_mask == preserved_mask + assert loss_mask != default_mask + assert generator.get_text_from_loss_mask(token_ids, loss_mask) == [ + "REASONING_1\n\n\nANSWER_1<|im_end|>\n", + "REASONING_2\n\n\nANSWER_2<|im_end|>\n", + ] def test_qwen3_and_qwen3_5_diverge_on_multi_turn_qwen35_data(): @@ -231,6 +333,49 @@ def test_qwen3_5_matches_expected_mask_for_tool_call_flow(): assert token_ids == expected_token_ids assert loss_mask == expected_mask assert generator.get_text_from_loss_mask(token_ids, loss_mask) == [ - "\n\n\nTOOL_CALL\n\n\n\n\nls\n\n\n<|im_end|>\n", + "TOOL_CALL\n\n\n\n\nls\n\n\n<|im_end|>\n", "REASONING\n\n\nFINAL<|im_end|>\n", ] + + +def test_qwen3_5_multimodal_alignment_inserts_media_mask_at_turn_position(): + tokenizer = FakeQwen35Tokenizer() + generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type="qwen3_5") + messages = [ + {"role": "user", "content": "USER_1"}, + {"role": "assistant", "content": "ANSWER_1"}, + { + "role": "user", + "content": [ + {"type": "image", "image": "image.jpg"}, + {"type": "text", "text": "USER_2"}, + ], + }, + {"role": "assistant", "content": "ANSWER_2"}, + ] + rendered_text, rendered_mask = tokenizer.render_with_expected_mask(messages) + image_token = "<|image_pad|>" + image_pos = rendered_text.index(image_token) + expanded_image_tokens = 3 + expanded_rendered_text = rendered_text.replace(image_token, image_token * expanded_image_tokens, 1) + input_ids = tokenizer(expanded_rendered_text, add_special_tokens=False)["input_ids"] + expected_mask = ( + rendered_mask[:image_pos] + + [0] * len(image_token) * expanded_image_tokens + + rendered_mask[image_pos + len(image_token) :] + ) + + _, loss_mask = generator.get_loss_mask_with_multimodal_alignment( + messages, + input_ids, + rendered_text=rendered_text, + image_grid_thw=[[1, 3, 4]], + image_token=image_token, + image_merge_size=2, + ) + + assert loss_mask == expected_mask + assert sum(loss_mask[image_pos : image_pos + len(image_token) * expanded_image_tokens]) == 0 + assert loss_mask[expanded_rendered_text.index("ANSWER_1")] == 1 + assert loss_mask[expanded_rendered_text.index("ANSWER_2")] == 1 + assert "<|im_start|>user" not in "".join(generator.get_text_from_loss_mask(input_ids, loss_mask)) diff --git a/tools/convert_torch_dist_to_hf_bridge.py b/tools/convert_torch_dist_to_hf_bridge.py index 798503f218..09715c2a35 100644 --- a/tools/convert_torch_dist_to_hf_bridge.py +++ b/tools/convert_torch_dist_to_hf_bridge.py @@ -1,12 +1,95 @@ import argparse +from dataclasses import dataclass +import io import os +from typing import Optional -import megatron.bridge.training.model_load_save as _model_load_save_module -from megatron.bridge import AutoBridge +import torch from slime.utils.megatron_bridge_utils import patch_auto_bridge_hf_config +def _patch_hybrid_pattern_compat(): + """Patch Megatron-LM forks that predate unified hybrid/MTP pattern helpers. + + The Slime v0.3.0 image pairs a newer Megatron-Bridge with a Megatron-LM fork + whose mamba_hybrid_layer_allocation module does not expose + parse_hybrid_pattern. Bridge imports the symbol unconditionally while loading + MLM checkpoints, even for non-Mamba models such as Qwen3.6-VL. + """ + import megatron.core.ssm.mamba_hybrid_layer_allocation as hybrid_alloc + + symbols = hybrid_alloc.Symbols + if not hasattr(symbols, "MTP_SEPARATOR"): + symbols.MTP_SEPARATOR = "/" + if not hasattr(symbols, "PIPE"): + symbols.PIPE = "|" + if not hasattr(symbols, "VALID_LAYERS"): + symbols.VALID_LAYERS = getattr(symbols, "VALID", None) + + if hasattr(hybrid_alloc, "parse_hybrid_pattern"): + return + + @dataclass + class ParsedHybridPattern: + main_pattern: Optional[str] + mtp_pattern: Optional[str] + mtp_num_depths: int + + def _valid_layer_symbols(): + valid = getattr(symbols, "VALID_LAYERS", None) or getattr(symbols, "VALID", None) + if valid is None: + valid = {symbols.MAMBA, symbols.ATTENTION, symbols.MLP, symbols.MOE} + return set(valid) + + def _validate_pattern(pattern: Optional[str], pattern_name: str) -> None: + if not pattern: + return + valid = _valid_layer_symbols() + for symbol in pattern: + if symbol == symbols.PIPE: + continue + if symbol not in valid: + raise ValueError(f"In {pattern_name} hybrid pattern, '{symbol}' is not one of {valid}") + + def parse_hybrid_pattern(pattern: Optional[str]) -> ParsedHybridPattern: + if not pattern: + return ParsedHybridPattern(main_pattern=None, mtp_pattern=None, mtp_num_depths=0) + + parts = pattern.split(symbols.MTP_SEPARATOR) + main_pattern = parts[0] or None + _validate_pattern(main_pattern, "main") + + if len(parts) == 1: + return ParsedHybridPattern(main_pattern=main_pattern, mtp_pattern=None, mtp_num_depths=0) + + mtp_patterns = parts[1:] + first_mtp_pattern = mtp_patterns[0] + if not first_mtp_pattern: + raise ValueError(f"Empty MTP hybrid pattern in '{pattern}'") + for mtp_pattern in mtp_patterns: + if mtp_pattern != first_mtp_pattern: + raise ValueError(f"Inconsistent MTP hybrid patterns in '{pattern}'") + _validate_pattern(mtp_pattern, "MTP") + + return ParsedHybridPattern( + main_pattern=main_pattern, + mtp_pattern=first_mtp_pattern, + mtp_num_depths=len(mtp_patterns), + ) + + hybrid_alloc.ParsedHybridPattern = ParsedHybridPattern + hybrid_alloc.parse_hybrid_pattern = parse_hybrid_pattern + + +_patch_hybrid_pattern_compat() + +import megatron.bridge.training.model_load_save as _model_load_save_module +from megatron.bridge import AutoBridge +from megatron.bridge.models.conversion import model_bridge as _model_bridge_module +from megatron.core.dist_checkpointing.strategies import torch as _mcore_torch_strategy + + # Here we need to patch Megatron Bridge's `load_model_config`, since the checkpoint is saved # by Megatron and lack of provider information. _provider_override = {} @@ -28,6 +111,57 @@ def _patched_load_model_config(checkpoint_path): _model_load_save_module.load_model_config = _patched_load_model_config +_bridge_class = getattr(_model_bridge_module, "ModelBridge", None) +if _bridge_class is None: + _bridge_class = getattr(_model_bridge_module, "MegatronModelBridge") +_original_build_conversion_tasks = _bridge_class.build_conversion_tasks + + +def _patched_build_conversion_tasks(self, *args, **kwargs): + tasks = list(_original_build_conversion_tasks(self, *args, **kwargs)) + filtered = [task for task in tasks if task is not None] + skipped = len(tasks) - len(filtered) + if skipped: + print(f"[convert] Skipping {skipped} unmapped Megatron->HF conversion tasks.") + return filtered + + +_bridge_class.build_conversion_tasks = _patched_build_conversion_tasks + + +_original_replace_sharded_keys = _mcore_torch_strategy._replace_sharded_keys_with_state_dict_keys + + +def _patched_replace_sharded_keys_with_state_dict_keys(state_dict, flat_mapping, rename_mapping): + """Handle ShardedObject BytesIO values when loading Slime torch_dist ckpts. + + This Megatron-LM fork unwraps tensor shards to lists, but leaves sharded + object payloads as a single BytesIO. The original restore function then + calls len(BytesIO) and fails. Decode those payloads back into the list that + mcore_to_pyt_state_dict serialized, and leave tensor shard lists unchanged. + """ + recovered_sd = {} + for key, tensors in state_dict.items(): + if isinstance(tensors, io.BytesIO): + tensors.seek(0) + tensors = torch.load(tensors, map_location="cpu", weights_only=False) + elif not isinstance(tensors, list): + tensors = [tensors] + + assert len(tensors) == len(rename_mapping[key]), ( + key, + len(tensors), + len(rename_mapping[key]), + ) + for tensor, recovered_key in zip(tensors, rename_mapping[key]): + recovered_sd[recovered_key] = tensor + + return _mcore_torch_strategy.unflatten_state_dict(recovered_sd, flat_mapping) + + +_mcore_torch_strategy._replace_sharded_keys_with_state_dict_keys = _patched_replace_sharded_keys_with_state_dict_keys + + if __name__ == "__main__": parser = argparse.ArgumentParser( description="Convert torch distributed checkpoint to HuggingFace format using Megatron Bridge" diff --git a/train_async.py b/train_async.py index 9d4c9b6473..d72ec8968f 100644 --- a/train_async.py +++ b/train_async.py @@ -1,3 +1,5 @@ +import logging + import ray from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models @@ -5,6 +7,30 @@ from slime.utils.logging_utils import configure_logger, finish_tracking, init_tracking from slime.utils.misc import should_run_periodic_action +logger = logging.getLogger(__name__) + + +def _free_rollout_data_refs(rollout_data_ref): + """Release Ray object-store copies of rollout data before checkpoint save.""" + if rollout_data_ref is None: + return + + object_refs = [] + for shard in rollout_data_ref: + inner = getattr(shard, "inner", None) + if isinstance(inner, ray.ObjectRef): + object_refs.append(inner) + + if not object_refs: + return + + try: + from ray._private.internal_api import free + + free(object_refs, local_only=False) + except Exception as exc: + logger.warning("Failed to free rollout Ray object refs before checkpoint save: %s", exc) + # The framework supports other asynchronous approaches such as fully async (which is shown in examples/full_async). def train(args): @@ -34,9 +60,15 @@ def train(args): if rollout_data_next_future is not None: rollout_data_curr_ref = ray.get(rollout_data_next_future) - # Start the next rollout early. - if rollout_id + 1 < args.num_rollout: + will_save = should_run_periodic_action(rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout) + + # Start the next rollout early, except before checkpoint save. Saving a large + # model already creates a high host-memory peak, so avoid holding the next + # rollout batch in Ray object store at the same time. + if rollout_id + 1 < args.num_rollout and not will_save: rollout_data_next_future = rollout_manager.generate.remote(rollout_id + 1) + else: + rollout_data_next_future = None if args.use_critic: actor_trains_this_step = rollout_id >= args.num_critic_only_steps @@ -48,7 +80,14 @@ def train(args): else: ray.get(actor_model.async_train(rollout_id, rollout_data_curr_ref)) - if should_run_periodic_action(rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout): + if will_save: + _free_rollout_data_refs(rollout_data_curr_ref) + rollout_data_curr_ref = None + if (not args.use_critic) or rollout_id >= args.num_critic_only_steps: + actor_model.clear_memory() + if args.use_critic: + critic_model.clear_memory() + if (not args.use_critic) or rollout_id >= args.num_critic_only_steps: actor_model.save_model( rollout_id, @@ -62,6 +101,9 @@ def train(args): if args.rollout_global_dataset: ray.get(rollout_manager.save.remote(rollout_id)) + if rollout_id + 1 < args.num_rollout: + rollout_data_next_future = rollout_manager.generate.remote(rollout_id + 1) + if (rollout_id + 1) % args.update_weights_interval == 0: # sync generate before update weights to prevent update weight in the middle of generation rollout_data_curr_ref = ray.get(x) if (x := rollout_data_next_future) is not None else None