support multimodal qwen36 sft

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"

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

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:

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"

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)]