FP8 RL Training and Rollout Weight Synchronization

skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py

@@ -20,7 +20,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import gc
 from typing import Any, Dict, List, Optional, Union
 
 import torch
@@ -210,8 +209,6 @@ def offload_megatron_grads_to_cpu(models):
             for _, param in model_chunk.named_parameters():
                 if param.grad is not None:
                     param.grad = param.grad.to("cpu", non_blocking=True)
-    gc.collect()
-    torch.cuda.empty_cache()
 
 
 @torch.no_grad()
@@ -223,8 +220,6 @@ def load_megatron_grads_to_gpu(models):
             for _, param in model_chunk.named_parameters():
                 if param.grad is not None:
                     param.grad = param.grad.to(torch.cuda.current_device(), non_blocking=True)
-    gc.collect()
-    torch.cuda.empty_cache()
 
 
 @torch.no_grad()
@@ -259,8 +254,6 @@ def offload_megatron_model_to_cpu(models):
         else:
             for _, param in model_chunk.named_parameters():
                 param.data = param.data.to("cpu", non_blocking=True)
-    gc.collect()
-    torch.cuda.empty_cache()
 
 
 @torch.no_grad()
@@ -280,8 +273,6 @@ def load_megatron_model_to_gpu(models):
             device_id = torch.cuda.current_device()
             for _, param in model_chunk.named_parameters():
                 param.data = param.data.to(device_id, non_blocking=True)
-    gc.collect()
-    torch.cuda.empty_cache()
 
 
 @torch.no_grad()
@@ -381,8 +372,6 @@ def _iter_opts(opt):
                 v["exp_avg"] = v["exp_avg"].to("cpu", non_blocking=True)
             if "exp_avg_sq" in v:
                 v["exp_avg_sq"] = v["exp_avg_sq"].to("cpu", non_blocking=True)
-        gc.collect()
-        torch.cuda.empty_cache()
 
 
 @torch.no_grad()
@@ -404,8 +393,6 @@ def _iter_opts(opt):
                     v["exp_avg"] = v["exp_avg"].to(torch.cuda.current_device(), non_blocking=True)
                 if "exp_avg_sq" in v:
                     v["exp_avg_sq"] = v["exp_avg_sq"].to(torch.cuda.current_device(), non_blocking=True)
-        gc.collect()
-        torch.cuda.empty_cache()
 
 
 def preprocess_packed_seqs(
@@ -426,12 +413,11 @@ def preprocess_packed_seqs(
       per row. This is the historical SkyRL behavior used by the RL path
       and the existing SFT path without mini-batch packing.
     - ``sub_seq_lengths is not None``: each row may contain multiple
-      sub-sequences concatenated end-to-end. ``sub_seq_lengths[r]`` lists
-      the per-sub-sequence valid token counts for row ``r``. Tokens
-      ``input_ids[r, :sum(sub_seq_lengths[r])]`` are assumed to be the
-      concatenated sub-sequences in order; any trailing tokens in the row
-      are pad. ``cu_seqlens`` enumerates every sub-sequence across every
-      row.
+      sub-sequences. ``sub_seq_lengths[r]`` lists their valid token counts.
+      Each sub-sequence begins at the next ``align_size`` boundary, so internal
+      alignment padding may separate adjacent sub-sequences; any remaining
+      trailing tokens are pad. ``cu_seqlens`` enumerates every sub-sequence
+      across every row.
 
     CP splits sequence into CP*2 chunks, and each GPU gets 2 chunks (GPU0
     gets first and last chunks, GPU1 gets second and second last chunks,

skyrl/backends/skyrl_train/distributed/megatron/model_utils.py

@@ -17,6 +17,7 @@
 # limitations under the License.
 
 import warnings
+from math import prod
 from typing import Any, Optional
 
 import megatron.core.parallel_state as mpu
@@ -1001,6 +1002,8 @@ def vocab_parallel_entropy_packed_sequences(
     loss_mask: Optional[torch.Tensor],
     cp_group: Optional[torch.distributed.ProcessGroup],
     sub_seq_lengths: Optional[list[list[int]]] = None,
+    chunk_size: Optional[int] = None,
+    chunk_memory_mb: int = 512,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Compute action-token entropy directly on TP+CP sharded packed logits.
 
@@ -1009,9 +1012,8 @@ def vocab_parallel_entropy_packed_sequences(
         local term is normalized by the global action-token count. Megatron's
         schedule already applies the CP loss scale for two-output loss funcs.
     """
-    entropy_tokens = vocab_parallel_entropy(vocab_parallel_logits).squeeze(0)
-    device = entropy_tokens.device
-    dtype = entropy_tokens.dtype
+    device = vocab_parallel_logits.device
+    dtype = vocab_parallel_logits.dtype
 
     attention_mask = attention_mask.to(device=device, dtype=torch.bool)
     cu_seqlens_padded = cu_seqlens_padded.to(device=device, dtype=torch.long)
@@ -1055,13 +1057,18 @@ def vocab_parallel_entropy_packed_sequences(
         cp_rank_for_token, local_indices = _packed_cp_rank_and_local_indices(
             cu_seqlens_padded, seq_indices, seq_offsets, seq_lens_padded, cp_size
         )
-        local_weights = torch.zeros_like(entropy_tokens)
+        local_weights = torch.zeros((int(vocab_parallel_logits.shape[-2]),), dtype=dtype, device=device)
         current_rank_mask = cp_rank_for_token == cp_rank
         local_weights[local_indices[current_rank_mask]] = packed_weights[current_rank_mask]
     else:
         local_weights = packed_weights
 
-    local_entropy_sum = (entropy_tokens * local_weights).sum()
+    local_entropy_sum = vocab_parallel_entropy_weighted_sum(
+        vocab_parallel_logits,
+        local_weights,
+        chunk_size=chunk_size,
+        chunk_memory_mb=chunk_memory_mb,
+    )
     local_count = local_weights.sum()
     global_count = local_count.detach().clone()
     global_entropy_sum = local_entropy_sum.detach().clone()
@@ -1333,7 +1340,51 @@ def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
         return softmax_logits
 
 
-def vocab_parallel_entropy(vocab_parallel_logits: torch.Tensor) -> torch.Tensor:
+def _floor_power_of_two(value: int) -> int:
+    return 1 << (value.bit_length() - 1)
+
+
+def _resolve_vocab_entropy_chunk_size(
+    vocab_parallel_logits: torch.Tensor,
+    chunk_size: Optional[int],
+    chunk_memory_mb: int,
+    peak_factor: int = 4,
+) -> Optional[int]:
+    """Resolve the sequence chunk size for vocab entropy.
+
+    ``None`` disables chunking. ``0`` uses the runtime vocab shard and dtype to
+    choose a size. Positive values specify the number of tokens per chunk.
+    """
+    if chunk_size is None:
+        return None
+    if chunk_size < 0:
+        raise ValueError(f"chunk_size must be non-negative or None, got {chunk_size}")
+    if chunk_memory_mb <= 0:
+        raise ValueError(f"chunk_memory_mb must be positive, got {chunk_memory_mb}")
+    seq_len = int(vocab_parallel_logits.shape[-2])
+    if seq_len <= 0:
+        return None
+    if chunk_size > 0:
+        return chunk_size if chunk_size < seq_len else None
+
+    budget_bytes = int(chunk_memory_mb) * 1024 * 1024
+    leading_elements = prod(vocab_parallel_logits.shape[:-2])
+    bytes_per_token = (
+        leading_elements * int(vocab_parallel_logits.shape[-1]) * vocab_parallel_logits.element_size() * peak_factor
+    )
+    if bytes_per_token <= 0:
+        return None
+
+    auto_chunk = max(1, min(seq_len, budget_bytes // bytes_per_token))
+    auto_chunk = _floor_power_of_two(auto_chunk)
+    return auto_chunk if auto_chunk < seq_len else None
+
+
+def vocab_parallel_entropy(
+    vocab_parallel_logits: torch.Tensor,
+    chunk_size: Optional[int] = None,
+    chunk_memory_mb: int = 512,
+) -> torch.Tensor:
     """Compute entropy when the logits are sharded in tp ranks
 
     Args:
@@ -1342,4 +1393,38 @@ def vocab_parallel_entropy(vocab_parallel_logits: torch.Tensor) -> torch.Tensor:
     Returns: (total_nnz,)
 
     """
-    return _VocabParallelEntropy.apply(vocab_parallel_logits)
+    resolved_chunk_size = _resolve_vocab_entropy_chunk_size(vocab_parallel_logits, chunk_size, chunk_memory_mb)
+    if resolved_chunk_size is None:
+        return _VocabParallelEntropy.apply(vocab_parallel_logits)
+
+    entropy_chunks = []
+    seq_len = int(vocab_parallel_logits.shape[-2])
+    for start in range(0, seq_len, resolved_chunk_size):
+        end = min(start + resolved_chunk_size, seq_len)
+        entropy_chunks.append(_VocabParallelEntropy.apply(vocab_parallel_logits[..., start:end, :]))
+    return torch.cat(entropy_chunks, dim=-1)
+
+
+def vocab_parallel_entropy_weighted_sum(
+    vocab_parallel_logits: torch.Tensor,
+    weights: torch.Tensor,
+    chunk_size: Optional[int] = None,
+    chunk_memory_mb: int = 512,
+) -> torch.Tensor:
+    """Compute ``sum(entropy * weights)`` with bounded temporary memory."""
+    resolved_chunk_size = _resolve_vocab_entropy_chunk_size(vocab_parallel_logits, chunk_size, chunk_memory_mb)
+    if resolved_chunk_size is None:
+        entropy_tokens = _VocabParallelEntropy.apply(vocab_parallel_logits)
+        return (entropy_tokens * weights).sum()
+
+    # Keep an autograd edge even when every chunk is masked out.
+    local_entropy_sum = vocab_parallel_logits[..., :0, :].sum()
+    seq_len = int(vocab_parallel_logits.shape[-2])
+    for start in range(0, seq_len, resolved_chunk_size):
+        end = min(start + resolved_chunk_size, seq_len)
+        weight_chunk = weights[start:end]
+        if torch.count_nonzero(weight_chunk).item() == 0:
+            continue
+        entropy_chunk = _VocabParallelEntropy.apply(vocab_parallel_logits[..., start:end, :])
+        local_entropy_sum = local_entropy_sum + (entropy_chunk * weight_chunk).sum()
+    return local_entropy_sum

skyrl/backends/skyrl_train/distributed/megatron/packing_utils.py

@@ -5,23 +5,29 @@
 def is_fp8_enabled(fp8: Any) -> bool:
     """Return whether a Megatron/TE fp8 config value enables FP8 execution."""
     if isinstance(fp8, str):
-        return fp8.strip().lower() not in {"", "0", "false", "none", "null", "no"}
+        return fp8.strip().lower() not in {"", "0", "false", "none", "null", "no", "off"}
     return bool(fp8)
 
 
 def get_packed_seq_align_size(tp_size: int, cp_size: int, fp8_enabled: bool = False) -> int:
-    """Return global per-subsequence padding needed for TP/CP layout."""
+    """Return the global alignment unit for packed TP/CP/FP8 sequences."""
+    if tp_size < 1 or cp_size < 1:
+        raise ValueError(f"tp_size and cp_size must be positive, got tp_size={tp_size}, cp_size={cp_size}")
     if cp_size > 1:
         layout_align = tp_size * cp_size * 2
     else:
         layout_align = tp_size
     if not fp8_enabled:
         return layout_align
-    return math.lcm(layout_align, 16 * cp_size)
+    fp8_token_align = 128 * tp_size * cp_size if tp_size > 1 else 16 * cp_size
+    return math.lcm(layout_align, fp8_token_align)
 
 
 def get_unpacked_seq_align_size(tp_size: int, fp8_enabled: bool = False) -> int:
-    """Return sequence padding needed when removing microbatch padding without CP."""
+    """Return the alignment unit for unpacked TP/FP8 sequences without CP."""
+    if tp_size < 1:
+        raise ValueError(f"tp_size must be positive, got {tp_size}")
     if not fp8_enabled:
         return tp_size
-    return math.lcm(tp_size, 16)
+    fp8_token_align = 128 * tp_size if tp_size > 1 else 16
+    return math.lcm(tp_size, fp8_token_align)