diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index ffdf801f684..7a72102ce20 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -401,3 +401,19 @@ logger: cluster: gpus_per_node: 1 num_nodes: 1 + +# TransferQueue-mediated data plane for sync GRPO. +# Off by default ā€” the legacy grpo_train trainer never engages this. +# Flip enabled=true and run grpo_train_sync to use TQ-mediated bulk +# transfer between rollout and train. See nemo_rl/data_plane/README.md. +data_plane: + enabled: false + impl: transfer_queue + backend: "simple" # TQ storage backend ('simple' or 'mooncake_cpu') + storage_capacity: 1000000 # max samples retained per partition + num_storage_units: 2 # storage shards + claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence + global_segment_size: 549755813888 # 512 GiB ā€” used when backend == "mooncake_cpu" + local_buffer_size: 68719476736 # 64 GiB ā€” used when backend == "mooncake_cpu" + # observability: # NotRequired + # enabled: false diff --git a/examples/configs/recipes/llm/grpo-glm47-flash-4n8g-automodel.yaml b/examples/configs/recipes/llm/grpo-glm47-flash-4n8g-automodel.yaml index ef7dcfc514f..3af609740eb 100644 --- a/examples/configs/recipes/llm/grpo-glm47-flash-4n8g-automodel.yaml +++ b/examples/configs/recipes/llm/grpo-glm47-flash-4n8g-automodel.yaml @@ -7,6 +7,7 @@ loss_fn: reference_policy_kl_penalty: 0.0 use_importance_sampling_correction: true truncated_importance_sampling_ratio: 2 + truncated_importance_sampling_type: tis checkpointing: checkpoint_dir: results/grpo-glm47-flash-4n8g-automodel policy: diff --git a/examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-dapo-4n8g-automodel.yaml b/examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-dapo-4n8g-automodel.yaml index 0fca436303b..09bcf82d7b3 100644 --- a/examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-dapo-4n8g-automodel.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-dapo-4n8g-automodel.yaml @@ -12,6 +12,7 @@ loss_fn: reference_policy_kl_penalty: 0.0 use_importance_sampling_correction: true truncated_importance_sampling_ratio: 2 + truncated_importance_sampling_type: tis ratio_clip_max: 0.28 ratio_clip_c: 10 checkpointing: diff --git a/examples/configs/recipes/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16.yaml b/examples/configs/recipes/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16.yaml index f9bd24b2244..2cf95a5f631 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16.yaml @@ -5,6 +5,7 @@ loss_fn: reference_policy_kl_penalty: 0.0 use_importance_sampling_correction: true truncated_importance_sampling_ratio: 2 + truncated_importance_sampling_type: tis checkpointing: checkpoint_dir: results/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16 policy: diff --git a/examples/configs/recipes/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-megatron-ep16.yaml b/examples/configs/recipes/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-megatron-ep16.yaml index a62b18017f5..b414e7dad35 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-megatron-ep16.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-megatron-ep16.yaml @@ -5,6 +5,7 @@ loss_fn: reference_policy_kl_penalty: 0.0 use_importance_sampling_correction: true truncated_importance_sampling_ratio: 2 + truncated_importance_sampling_type: tis checkpointing: checkpoint_dir: results/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-megatron-ep16 policy: diff --git a/examples/run_grpo.py b/examples/run_grpo.py index 9694fa396c6..d1b9ed23b6e 100644 --- a/examples/run_grpo.py +++ b/examples/run_grpo.py @@ -31,6 +31,22 @@ from nemo_rl.utils.logger import get_next_experiment_dir +def _select_trainer(master_config: MasterConfig): + """Pick the synchronous trainer based on ``data_plane.enabled``. + + Factored out so test_architecture_invariants can verify dispatch + without the full setup() path. + """ + dp_cfg = master_config.data_plane or {} + if dp_cfg.get("enabled", False): + from nemo_rl.algorithms.grpo_sync import grpo_train_sync + + print("šŸš€ Running synchronous GRPO training (TransferQueue)") + return grpo_train_sync + print("šŸš€ Running synchronous GRPO training (legacy)") + return grpo_train + + def parse_args() -> tuple[argparse.Namespace, list[str]]: """Parse command line arguments.""" parser = argparse.ArgumentParser(description="Run GRPO training with configuration") @@ -100,6 +116,20 @@ def main() -> None: val_task_to_env, ) = setup_response_data(tokenizer, config.data, config.env) + # Pick the policy factory at the launcher level so the legacy trainer + # stays data-plane-agnostic (architectural invariant ā€” see + # tests/data_plane/unit/test_architecture_invariants.py). + _dp_cfg = config.data_plane or {} + if _dp_cfg.get("enabled", False): + from nemo_rl.models.policy.tq_policy import TQPolicy + + def _make_policy(**kwargs): + return TQPolicy(**kwargs, dp_cfg=_dp_cfg) + + _policy_factory = _make_policy + else: + _policy_factory = None # setup() defaults to plain Policy + ( policy, policy_generation, @@ -111,7 +141,13 @@ def main() -> None: checkpointer, grpo_state, master_config, - ) = setup(config, tokenizer, dataset, val_dataset) + ) = setup( + config, + tokenizer, + dataset, + val_dataset, + policy_factory=_policy_factory, + ) # Check if async mode is enabled if "async_grpo" in config.grpo and config.grpo["async_grpo"]["enabled"]: @@ -165,10 +201,10 @@ def main() -> None: max_trajectory_age_steps=async_config["max_trajectory_age_steps"], ) else: - print("šŸš€ Running synchronous GRPO training") - - # Run standard GRPO training - grpo_train( + # Two parallel synchronous trainers (verl-style ā€” main_ppo.py vs + # main_ppo_sync.py). data_plane.enabled selects which one runs. + trainer = _select_trainer(master_config) + trainer( policy, policy_generation, dataloader, diff --git a/nemo_rl/algorithms/distillation.py b/nemo_rl/algorithms/distillation.py index 76179f7c8b6..5d49638051c 100644 --- a/nemo_rl/algorithms/distillation.py +++ b/nemo_rl/algorithms/distillation.py @@ -530,7 +530,7 @@ def distillation_train( student_generation = student_policy # type: ignore NEED_REFIT = False POLICY_GENERATION_STALE = True # tracks if generation needs a refit before running - assert student_generation is not None # for mypy type check + assert student_generation is not None # common config/state items current_epoch = distillation_save_state["current_epoch"] # current epoch diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index f0731aaccc5..a1e7b69dd17 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -17,7 +17,7 @@ import warnings from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext -from typing import Any, NotRequired, Optional, TypedDict, TypeVar, cast +from typing import Any, Callable, NotRequired, Optional, TypedDict, TypeVar, cast import numpy as np import ray @@ -59,6 +59,7 @@ get_keys_from_message_log, ) from nemo_rl.data.utils import extract_necessary_env_names, load_dataloader_state +from nemo_rl.data_plane.interfaces import DataPlaneConfig from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env from nemo_rl.distributed.virtual_cluster import ClusterConfig, RayVirtualCluster @@ -207,6 +208,7 @@ class MasterConfig(BaseModel, extra="allow"): logger: GRPOLoggerConfig cluster: ClusterConfig checkpointing: CheckpointingConfig + data_plane: Optional[DataPlaneConfig] = None # =============================================================================== @@ -220,6 +222,7 @@ def setup( dataset: AllTaskProcessedDataset | dict[str, AllTaskProcessedDataset], val_dataset: Optional[AllTaskProcessedDataset], processor: Optional[AutoProcessor] = None, + policy_factory: Optional[Callable[..., ColocatablePolicyInterface]] = None, ) -> tuple[ ColocatablePolicyInterface, Optional[GenerationInterface], @@ -580,10 +583,15 @@ def init_train_dataloader(dataset, suffix: str = ""): "(reference model is not loaded)." ) + # Caller-supplied factory lets the sync trainer swap in a TQ-mediated + # Policy subclass without this shared setup needing to know the data + # plane exists. Default is the plain Policy class ā€” legacy behavior. + _make_policy = policy_factory if policy_factory is not None else Policy + def init_policy(): """Initialize policy training workers.""" t0 = time.perf_counter() - p = Policy( + p = _make_policy( cluster=train_cluster, config=policy_config, tokenizer=tokenizer, @@ -1360,7 +1368,7 @@ def grpo_train( policy_generation = policy # type: ignore NEED_REFIT = False POLICY_GENERATION_STALE = True # tracks if generation needs a refit before running - assert policy_generation is not None # for mypy type check + assert policy_generation is not None # Check if we need to sync KV cache scales # When fallback to policy as the policy_generation, we use getattr to check. diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py new file mode 100644 index 00000000000..c147da2cd45 --- /dev/null +++ b/nemo_rl/algorithms/grpo_sync.py @@ -0,0 +1,1270 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""GRPO trainer ā€” TransferQueue-mediated path (sync). + +Sibling fork of ``nemo_rl.algorithms.grpo``. Each file has zero +internal branching on whether TQ is engaged; the example script +chooses one or the other based on ``data_plane.enabled``. + +Setup and helpers are re-imported from ``grpo``; the training loop body +is duplicated here so the per-step lifecycle hooks (register / seed-put +/ per-rank fetch / clear) can live in straight sequential code. +Validation is implemented locally as :func:`validate_sync` ā€” a +TQ-mediated sibling of :func:`nemo_rl.algorithms.grpo.validate` that +routes val rollouts through ``SyncRolloutActor.rollout_to_tq`` into a +per-batch ``"val"`` partition. + +Parity with the legacy path is verified by running the same config +against both entrypoints and diffing the wandb runs. +""" + +from __future__ import annotations + +import gc +import os +import warnings +from typing import TYPE_CHECKING, Any, Optional + +if TYPE_CHECKING: + from nemo_rl.models.policy.tq_policy import TQPolicy + +import numpy as np +import ray +import torch +from torchdata.stateful_dataloader import StatefulDataLoader + +# Re-imports from grpo so this file is a thin trainer-only fork. +from nemo_rl.algorithms.grpo import ( + GRPOSaveState, + MasterConfig, + _create_advantage_estimator, + _log_mixed_rewards_and_advantages_information, + _should_log_nemo_gym_responses, + _should_use_nemo_gym, + compute_and_apply_seq_logprob_error_masking, + refit_policy_generation, + scale_rewards, +) +from nemo_rl.algorithms.loss import ( + ClippedPGLossDataDict, +) +from nemo_rl.algorithms.loss.interfaces import LossFunction +from nemo_rl.algorithms.reward_functions import apply_reward_shaping +from nemo_rl.algorithms.utils import ( + calculate_baseline_and_std_per_prompt, + get_gdpo_reward_component_keys, + log_generation_metrics_to_wandb, + print_performance_metrics, +) +from nemo_rl.data.interfaces import DatumSpec +from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message +from nemo_rl.data_plane.interfaces import KVBatchMeta +from nemo_rl.data_plane.schema import DP_CALIB_INPUT_FIELDS +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.environments.interfaces import EnvironmentInterface +from nemo_rl.experience.sync_rollout_actor import SyncRolloutActor +from nemo_rl.models.generation.interfaces import GenerationInterface +from nemo_rl.models.policy.interfaces import ColocatablePolicyInterface +from nemo_rl.utils.checkpoint import CheckpointManager +from nemo_rl.utils.logger import Logger, print_message_log_samples +from nemo_rl.utils.memory_tracker import MemoryTracker +from nemo_rl.utils.nsys import maybe_gpu_profile_step +from nemo_rl.utils.timer import TimeoutChecker, Timer +from nemo_rl.utils.venvs import make_actor_runtime_env + +# ā”€ā”€ DAPO non-zero-std dynamic sampling, slice-only ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ +# Slice-only formulation of nemo_rl.algorithms.grpo.dynamic_sampling: filter +# on std != 0, accumulate survivors across iterations, slice on overflow. +# Bulk in TQ untouched except for clear_samples of dropped/discarded uids. + + +def _apply_dynamic_sampling( + *, + meta: KVBatchMeta, + driver_carry: BatchedDataDict, + pending_meta: Optional[KVBatchMeta], + pending_carry: Optional[BatchedDataDict], + pending_unfiltered_rewards: list[torch.Tensor], + train_prompts_size: int, + num_gen_batches: int, + max_gen_batches: int, + policy: "TQPolicy", +) -> tuple[ + Optional[KVBatchMeta], + Optional[BatchedDataDict], + list[torch.Tensor], + bool, + dict[str, Any], + Optional[torch.Tensor], +]: + """Process one dynamic-sampling iteration. + + Drops zero-std (filtered) keys, merges survivors into the running + pending cache, and reports whether the cache has reached + ``train_prompts_size``. When complete, the returned ``pending_*`` IS + the training batch. + + Args: + meta: This iteration's ``KVBatchMeta``. + driver_carry: Per-row driver-local tensors for this iteration + (rewards, masks, prompt_ids_for_adv, baseline/std, ā€¦). + pending_meta: Survivors accumulated from prior iterations. + pending_carry: ``driver_carry`` rows aligned to ``pending_meta``. + pending_unfiltered_rewards: All iterations' rewards pre-filter, + for legacy reward metric parity. + train_prompts_size: Target batch size. + num_gen_batches: Iteration counter (1-based). + max_gen_batches: Upper bound on iterations before raising. + policy: TQPolicy whose ``discard_samples`` is used to drop filtered keys. + + Returns: + ``(pending_meta, pending_carry, pending_rewards, is_complete, + ds_metrics, unfiltered_for_log)``. + """ + # Cumulative unfiltered total_reward for legacy metrics["reward"] + # parity. Reference-only append (no copy) ā€” slice tensors are + # produced fresh per iteration, not aliased to TQ-owned bulk. + pending_unfiltered_rewards.append(driver_carry["total_reward"]) + + # Filter input comes from ``meta.tags`` so the filter decision is + # meta-only ā€” no tensor data needed. The driver mirrored ``std`` + # into tags right after baseline/std compute. + if meta.tags is None: + raise ValueError( + "_apply_dynamic_sampling: meta.tags is None ā€” driver must " + "stamp 'std' into meta.tags before this call." + ) + keep_idx = [i for i, t in enumerate(meta.tags) if t["std"] != 0.0] + drop_keys = [k for k, t in zip(meta.sample_ids, meta.tags) if t["std"] == 0.0] + if drop_keys: + policy.discard_samples(drop_keys, meta.partition_id) + + # Subset survivors and merge into the running cache. + if keep_idx: + survivors_meta = meta.subset(keep_idx) + survivors_carry = driver_carry.select_indices(keep_idx) + survivors_carry["filtered_reward"] = survivors_carry["total_reward"] + if pending_meta is None: + pending_meta, pending_carry = survivors_meta, survivors_carry + else: + assert pending_carry is not None + pending_meta = pending_meta.concat(survivors_meta) + pending_carry = BatchedDataDict.from_batches( + [pending_carry, survivors_carry] + ) + + n = len(pending_meta.sample_ids) if pending_meta is not None else 0 + if n < train_prompts_size: + if num_gen_batches > max_gen_batches: + raise ValueError( + f"Dynamic sampling reached max_gen_batches={max_gen_batches}. " + f"Increase grpo.dynamic_sampling_max_gen_batches or revisit " + f"data diversity / num_prompts_per_step / num_generations_per_prompt." + ) + return pending_meta, pending_carry, pending_unfiltered_rewards, False, {}, None + + ds_metrics: dict[str, Any] = {"dynamic_sampling_num_gen_batches": num_gen_batches} + assert pending_meta is not None and pending_carry is not None + if n > train_prompts_size: + policy.discard_samples( + list(pending_meta.sample_ids[train_prompts_size:]), + pending_meta.partition_id, + ) + pending_meta = pending_meta.slice(0, train_prompts_size) + pending_carry = pending_carry.slice(0, train_prompts_size) + ds_metrics["dynamic_sampling_num_discarded_valid_samples"] = ( + n - train_prompts_size + ) + + unfiltered_for_log = torch.cat(pending_unfiltered_rewards)[:train_prompts_size] + return pending_meta, pending_carry, [], True, ds_metrics, unfiltered_for_log + + +def validate_sync( + *, + rollout_actor: SyncRolloutActor, + policy: "TQPolicy", + val_dataloader: Optional[StatefulDataLoader], + val_task_to_env: Optional[dict[str, EnvironmentInterface]], + step: int, + master_config: MasterConfig, + logger: Optional[Logger] = None, + partition_id: str = "val", +) -> tuple[dict[str, Any], dict[str, Any]]: + """TQ-mediated counterpart to :func:`nemo_rl.algorithms.grpo.validate`. + + Per-batch: register the val partition ā†’ ``rollout_to_tq`` ā†’ + ``policy.read_from_dataplane`` for message logs ā†’ ``policy.finish_step``. + Caller owns ``policy_generation.prepare_for_generation`` / + ``finish_generation`` around the call; the actor's per-rollout + ``finish_generation`` is suppressed so inference state stays warm + across batches. + """ + if val_dataloader is None: + assert master_config.grpo["val_period"] == 0, ( + "val_dataloader is None, so grpo.val_period must be 0" + ) + print(" āš ļø No validation dataloader provided, skipping validation", flush=True) + return {}, {} + + timer = Timer() + total_rewards: list[float] = [] + total_lengths: list[float] = [] + all_message_logs: list[list[dict[str, str]]] = [] + additional_metrics: dict[str, Any] = {} + capture_extras = _should_use_nemo_gym(master_config) + + with timer.time("total_validation_time"): + print(f"ā–¶ Starting validation at step {step}...", flush=True) + max_batches = ( + master_config.grpo["max_val_samples"] + // master_config.grpo["val_batch_size"] + ) + for batch_idx, val_batch in enumerate(val_dataloader): + if batch_idx >= max_batches: + break + n_prompts = int(val_batch.size) + policy.prepare_val_partition(n_prompts, partition_id=partition_id) + meta, driver_carry, rollout_metrics, _ = ray.get( + rollout_actor.rollout_to_tq.remote( + val_batch, + partition_id=partition_id, + first_iter=False, + finish_generation=False, + task_to_env_override=val_task_to_env, + carry_keys=["total_reward", "turn_roles", "turn_contents"], + ) + ) + roles = driver_carry["turn_roles"] + contents = driver_carry["turn_contents"] + total_rewards.extend(driver_carry["total_reward"].tolist()) + total_lengths.append(rollout_metrics["mean_gen_tokens_per_sample"]) + all_message_logs.extend( + [{"role": r, "content": c} for r, c in zip(roles[i], contents[i])] + for i in range(n_prompts) + ) + if capture_extras: + additional_metrics = rollout_metrics + policy.finish_step(meta) + + accuracy = ( + torch.tensor(total_rewards, dtype=torch.float32).mean().item() + if total_rewards + else 0.0 + ) + avg_length = sum(total_lengths) / len(total_lengths) if total_lengths else 0.0 + val_metrics = { + "accuracy": accuracy, + "avg_length": avg_length, + **additional_metrics, + } + try: + print_message_log_samples( + all_message_logs, + total_rewards, + num_samples=min( + master_config.logger["num_val_samples_to_print"], + len(all_message_logs), + ), + step=step, + ) + except Exception as e: + print(f"\n āš ļø Error displaying message samples: {str(e)}") + print(" āš ļø Continuing validation without displaying samples...", flush=True) + + timing_metrics = timer.get_timing_metrics(reduction_op="sum") + print( + f"\nšŸ“Š Validation Results:\n" + f" ā€¢ Accuracy: {accuracy:.4f}\n" + f" ā€¢ Average response length: {avg_length:.1f} tokens\n" + f" ā€¢ Samples processed: {len(total_rewards)}\n" + f" ā±ļø Total validation time: " + f"{timing_metrics.get('total_validation_time', 0):.2f}s", + flush=True, + ) + if logger is not None: + logger.log_batched_dict_as_jsonl( + {"content": all_message_logs, "rewards": total_rewards}, + f"val_data_step{step}.jsonl", + ) + timer.reset() + gc.collect() + torch.cuda.empty_cache() + return val_metrics, timing_metrics + + +def grpo_train_sync( + policy: ColocatablePolicyInterface, + policy_generation: Optional[GenerationInterface], + wrapped_dataloader, + val_dataloader: Optional[StatefulDataLoader], + tokenizer, + loss_fn: LossFunction, + task_to_env: dict[str, EnvironmentInterface], + val_task_to_env: Optional[dict[str, EnvironmentInterface]], + logger: Logger, + checkpointer: CheckpointManager, + grpo_save_state: GRPOSaveState, + master_config: MasterConfig, +) -> None: + """Run GRPO training algorithm ā€” TransferQueue-mediated. + + Body mirrors :func:`nemo_rl.algorithms.grpo.grpo_train` with TQ-mediated + Policy methods substituting the in-memory dispatch. The TQ lifecycle + (controller bootstrap, worker attach, partition register, fan-out, + drain, close) is fully encapsulated in + :class:`nemo_rl.models.policy.tq_policy.TQPolicy` ā€” this trainer just + calls ``policy.prepare_step``, ``policy.get_logprobs``, + ``policy.get_reference_policy_logprobs``, and ``policy.train``. + + Parity with the legacy path is verified by running the same config + against both entrypoints and diffing the wandb runs. + """ + timer = Timer() + timeout = TimeoutChecker( + timeout=master_config.checkpointing["checkpoint_must_save_by"], + fit_last_save_time=True, + ) + timeout.start_iterations() + memory_tracker = MemoryTracker() + + kv_scales_cache = None # Cache reused for computed kv scales + + NEED_REFIT = True + # If policy_generation is None, use the policy as the generation interface (megatron framework backend) + if policy_generation is None: + policy_generation = policy # type: ignore + NEED_REFIT = False + POLICY_GENERATION_STALE = True + assert policy_generation is not None + + if master_config.grpo.get("skip_reference_policy_logprobs_calculation"): + assert master_config.loss_fn.reference_policy_kl_penalty == 0 + print( + "Reference policy logprob calculation will be skipped since `grpo.skip_reference_policy_logprobs_calculation` is set to True and `loss_fn.reference_policy_kl_penalty` is 0." + ) + + sync_kv_scales = getattr(policy_generation, "requires_kv_scale_sync", False) + + current_step = grpo_save_state["current_step"] + total_steps = grpo_save_state["total_steps"] + max_num_steps = master_config.grpo["max_num_steps"] + current_epoch = grpo_save_state["current_epoch"] + max_num_epochs = master_config.grpo["max_num_epochs"] + consumed_samples = grpo_save_state["consumed_samples"] + total_valid_tokens = grpo_save_state.get("total_valid_tokens", 0) + val_at_start = master_config.grpo["val_at_start"] + val_at_end = master_config.grpo["val_at_end"] + val_period = master_config.grpo["val_period"] + colocated_inference = master_config.policy["generation"]["colocated"]["enabled"] + + adv_estimator = _create_advantage_estimator(master_config) + + # ā”€ā”€ Data-plane setup (mandatory in the sync trainer) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + # Sync trainer requires a TQ-mediated policy. The TQPolicy actor + # bootstraps the controller and attaches workers; ``policy.dp_cfg`` + # is the public marker. The explicit master_config check is the + # entry-guard so users running this trainer with the legacy policy + # see a clear error rather than an opaque AttributeError. + dp_cfg = master_config.data_plane + if not dp_cfg or not dp_cfg["enabled"]: + raise ValueError( + "grpo_train_sync requires master_config['data_plane']['enabled']=True. " + "Use the legacy nemo_rl.algorithms.grpo.grpo_train trainer if you don't " + "want TransferQueue." + ) + + # Driver-side pad-value dict for materialize() ā€” the wire emits + # jagged tensors for variable-length token fields (input_ids, + # prompt_ids_for_adv); other fields default to pad=0. + _pad_dict = { + "input_ids": tokenizer.pad_token_id, + "prompt_ids_for_adv": tokenizer.pad_token_id, + } + if not hasattr(policy, "dp_cfg"): + raise ValueError( + "grpo_train_sync requires a TQ-mediated policy " + "(nemo_rl.models.policy.tq_policy.TQPolicy). examples/run_grpo.py " + "constructs it via the policy_factory when data_plane.enabled=True." + ) + + # TQ-resident tensors live on CPU; baseline/std are computed on the + # slice without a CUDA hop. The flag is a no-op here ā€” warn so users + # don't expect it to do anything. + if master_config.grpo.get("calculate_advantages_on_gpu"): + warnings.warn( + "grpo.calculate_advantages_on_gpu has no effect when " + "data_plane.enabled=true; baseline/std are computed on CPU " + "because TQ-resident tensors are CPU-side.", + stacklevel=2, + ) + + # ā”€ā”€ Sync rollout actor (rollout 1-hop put) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + # The actor owns the multi-turn rollout loop AND post-rollout + # flatten / mask construction / prompt extraction / baseline-std / + # TQ first-write. Bulk tensors stay actor-side until put_samples; + # driver receives only KVBatchMeta + small slice via Ray. + rollout_actor = SyncRolloutActor.options( + runtime_env=make_actor_runtime_env( + "nemo_rl.experience.sync_rollout_actor.SyncRolloutActor" + ), + ).remote( + policy_generation=policy_generation, + tokenizer=tokenizer, + task_to_env=task_to_env, + master_config=master_config, + dp_cfg=dp_cfg, + ) + + if val_at_start and current_step == 0: + print("\nšŸ” Running initial validation...", flush=True) + memory_tracker.snapshot_start_of_stage("Initial validation", dir()) + + if NEED_REFIT and POLICY_GENERATION_STALE: + refit_policy_generation(policy, policy_generation, colocated_inference) + POLICY_GENERATION_STALE = False + else: + policy_generation.prepare_for_generation() + val_metrics, validation_timings = validate_sync( + rollout_actor=rollout_actor, + policy=policy, + val_dataloader=val_dataloader, + val_task_to_env=val_task_to_env, + step=0, + master_config=master_config, + logger=logger, + ) + policy_generation.finish_generation() + logger.log_metrics(val_metrics, current_step, prefix="validation") + logger.log_metrics(validation_timings, current_step, prefix="timing/validation") + + if master_config.data["use_multiple_dataloader"]: + warnings.warn( + "When using multiple dataloaders, MultipleDataloaderWrapper operates as an infinite iterator. " + "As a result, grpo.max_num_epochs will be ignored, and only grpo.max_num_steps will be used." + ) + + while current_epoch < max_num_epochs and total_steps < max_num_steps: + memory_tracker.snapshot_start_of_stage("Preparing batch", dir()) + print(f"\n{'=' * 25} Epoch {current_epoch + 1}/{max_num_epochs} {'=' * 25}") + # 1-hop cross-iteration cache for dynamic_sampling: across + # multiple inner iterations we accumulate non-zero-std prompts + # until we have enough for a full training batch. The TQ + # payload of pending uids remains alive until either consumed + # by training (clear_samples at step end) or evicted on overflow. + # ``pending_unfiltered_rewards`` is logging-only ā€” preserves + # legacy ``metrics["reward"]`` semantics (cumulative unfiltered + # total_reward across all contributing iterations). + pending_meta = None + pending_carry: Optional[BatchedDataDict] = None + pending_unfiltered_rewards: list[torch.Tensor] = [] + dynamic_sampling_num_gen_batches = 0 + + for batch in wrapped_dataloader: + metrics_logging_data: dict = {} + metrics: dict = {} + + if master_config.data["use_multiple_dataloader"]: + print( + f"\n{'=' * 25} Step {current_step + 1}/{max_num_steps} {'=' * 25}", + flush=True, + ) + else: + print( + f"\n{'=' * 25} Step {current_step + 1}/{min(len(wrapped_dataloader), max_num_steps)} {'=' * 25}", + flush=True, + ) + + maybe_gpu_profile_step(policy, total_steps + 1) + if policy != policy_generation: + maybe_gpu_profile_step(policy_generation, total_steps + 1) + val_metrics, validation_timings = None, None + + with timer.time("total_step_time"): + print("ā–¶ Preparing batch...", flush=True) + with timer.time("data_processing"): + repeated_batch: BatchedDataDict[DatumSpec] = ( + batch.repeat_interleave( + master_config.grpo["num_generations_per_prompt"] + ) + ) + + memory_tracker.snapshot_start_of_stage("Generation", dir()) + print( + f"ā–¶ Generating responses for batch of size {repeated_batch.size}...", + flush=True, + ) + with timer.time("prepare_for_generation/total"): + if NEED_REFIT and POLICY_GENERATION_STALE: + if sync_kv_scales and kv_scales_cache is None: + # KV-scale calibration uses message_log of the + # current step's PROMPTS (pre-generation), which + # is small and lives on the driver naturally. + # Unrelated to the rollout 1-hop put. + print("ā–¶ Computing KV cache scales...", flush=True) + policy.prepare_for_lp_inference() + calib_flat, calib_input_lengths = ( + batched_message_log_to_flat_message( + repeated_batch["message_log"], + pad_value_dict={ + "token_ids": tokenizer.pad_token_id + }, + make_sequence_length_divisible_by=master_config.policy[ + "make_sequence_length_divisible_by" + ], + ) + ) + calibration_data = BatchedDataDict[ClippedPGLossDataDict]( + { + "input_ids": calib_flat["token_ids"], + "input_lengths": calib_input_lengths, + } + ) + calibration_data.update( + calib_flat.get_multimodal_dict(as_tensors=False) + ) + calibration_data.to("cpu") + kv_scales_cache = policy.calibrate_qkv_fp8_scales( + calibration_data, include_q=True + )["layers"] + + refit_policy_generation( + policy, + policy_generation, + colocated_inference, + timer=timer, + kv_scales=kv_scales_cache if sync_kv_scales else None, + ) + POLICY_GENERATION_STALE = False + else: + if colocated_inference: + policy.offload_after_refit() + policy_generation.prepare_for_generation() + + # ā”€ā”€ Per-step TQ partition register ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + # Done before the rollout actor's put_samples so the + # partition exists with the expected schema. + policy.prepare_step( + num_samples=int(repeated_batch.size), + group_size=master_config.grpo["num_generations_per_prompt"], + ) + + # ā”€ā”€ Rollout 1-hop put: actor runs rollout + flatten + + # mask construction + prompt extraction + baseline/std, + # writes bulk to TQ in one flat put_samples, returns + # only meta + small slice. Bulk never visits the driver. + dynamic_sampling_num_gen_batches += 1 + with timer.time("generation"): + # Single Ray RPC: rollout + flatten + mask + prompt + # extraction + baseline/std + put_samples + finish + # generation + logger metrics ā€” all bundled into one + # round-trip. + # ``first_iter`` is the actor's signal to call + # ``policy_generation.snapshot_step_metrics()``. + # ``dynamic_sampling_num_gen_batches`` is incremented + # to 1 just above before this branch ā€” keep these in + # sync if either is renamed. + ( + meta, + driver_carry, + rollout_metrics, + generation_logger_metrics, + ) = ray.get( + rollout_actor.rollout_to_tq.remote( + repeated_batch, + partition_id=policy.tq_partition_id, + group_size=master_config.grpo["num_generations_per_prompt"], + first_iter=(dynamic_sampling_num_gen_batches == 1), + ) + ) + + if not _should_log_nemo_gym_responses(master_config): + for key in list(rollout_metrics): + if "full_result" in key: + rollout_metrics.pop(key) + + metrics_logging_data["mean_gen_tokens_per_sample"] = ( + rollout_metrics["mean_gen_tokens_per_sample"] + ) + logger.log_metrics(rollout_metrics, total_steps + 1, prefix="train") + + # ā”€ā”€ Per-sample driver compute on slice ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + # scale_rewards / apply_reward_shaping / overlong filter + # / baseline-std all operate on small per-sample + # tensors. Mirrors grpo_sync.py legacy layout ā€” they + # used to be on the driver, were briefly on the actor, + # now back on the driver where they belong (no bulk + # touched by any of these ops). + with timer.time("reward_calculation"): + driver_carry = scale_rewards( + driver_carry, + master_config.grpo["reward_scaling"], + ) + if master_config.grpo["reward_shaping"]["enabled"]: + driver_carry = apply_reward_shaping( + driver_carry, + master_config.grpo["reward_shaping"], + ) + driver_carry["baseline"], driver_carry["std"] = ( + calculate_baseline_and_std_per_prompt( + driver_carry["prompt_ids_for_adv"], + driver_carry["total_reward"], + torch.ones_like(driver_carry["total_reward"]), + leave_one_out_baseline=master_config.grpo[ + "use_leave_one_out_baseline" + ], + ) + ) + # Mirror std onto meta so dynamic_sampling can filter + # without fetching tensor data. + meta.stamp_tags( + { + "std": driver_carry["std"].tolist(), + "baseline": driver_carry["baseline"].tolist(), + } + ) + + # ā”€ā”€ Dynamic sampling (DAPO non-zero-std filter) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + # Slice-only; bulk in TQ untouched except for clear_samples + # of dropped / overflow-discarded uids. + ds_metrics: dict = {} + unfiltered_rewards_for_logging: Optional[torch.Tensor] = None + if master_config.grpo["use_dynamic_sampling"]: + with timer.time("dynamic_sampling"): + train_prompts_size = ( + master_config.grpo["num_prompts_per_step"] + * master_config.grpo["num_generations_per_prompt"] + ) + ( + pending_meta, + pending_carry, + pending_unfiltered_rewards, + is_complete, + ds_metrics, + unfiltered_rewards_for_logging, + ) = _apply_dynamic_sampling( + meta=meta, + driver_carry=driver_carry, + pending_meta=pending_meta, + pending_carry=pending_carry, + pending_unfiltered_rewards=pending_unfiltered_rewards, + train_prompts_size=train_prompts_size, + num_gen_batches=dynamic_sampling_num_gen_batches, + max_gen_batches=master_config.grpo[ + "dynamic_sampling_max_gen_batches" + ], + policy=policy, + ) + if not is_complete: + current_size = ( + len(pending_meta.sample_ids) + if pending_meta is not None + else 0 + ) + print( + f"Dynamic sampling: {current_size}/{train_prompts_size} " + f"non-zero-std prompts after batch " + f"{dynamic_sampling_num_gen_batches}; sampling more.", + flush=True, + ) + continue + + # Adopt the now-complete cache as this step's batch. + meta = pending_meta + driver_carry = pending_carry + pending_meta = None + pending_carry = None + + # Mirrors legacy ``grpo.py:1707-1716`` ā€” applied on the + # post-DS survivors so dropped rows don't affect this set. + if master_config.grpo["overlong_filtering"]: + lm = driver_carry["loss_multiplier"].clone() + lm[driver_carry["truncated"]] = 0 + driver_carry["loss_multiplier"] = lm + + # ā”€ā”€ Unpack slice (small per-sample tensors) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + rewards = ( + driver_carry["filtered_reward"] + if master_config.grpo["use_dynamic_sampling"] + else driver_carry["total_reward"] + ) + baseline = driver_carry["baseline"] + std = driver_carry["std"] + input_lengths = driver_carry["input_lengths"] + prompt_ids_for_adv = driver_carry["prompt_ids_for_adv"] + loss_multiplier = driver_carry["loss_multiplier"] + truncated = driver_carry["truncated"] + length = driver_carry["length"] + + gen_step_metrics = {} + if hasattr(policy_generation, "get_step_metrics"): + gen_step_metrics = policy_generation.get_step_metrics() + baseline_for_log = baseline.clone() + + memory_tracker.snapshot_start_of_stage("Computing logprobs", dir()) + print("ā–¶ Preparing for logprob inference...", flush=True) + with timer.time("logprob_inference_prep"): + policy.prepare_for_lp_inference() + + print("ā–¶ Computing logprobs...", flush=True) + with timer.time("policy_and_reference_logprobs"): + # Meta-driven worker dispatch. Workers fetch their + # slice from TQ and write ``prev_logprobs`` / + # ``reference_policy_logprobs`` columns back to TQ + # under ``meta.sample_ids``. The Ray return is + # discarded ā€” driver reads from TQ below in one + # batched fetch to avoid double-shipping the per-token + # tensor through Ray's plasma store on top of the TQ + # writeback. + policy.get_logprobs_from_meta(meta, timer=timer) + compute_ref = not master_config.grpo.get( + "skip_reference_policy_logprobs_calculation" + ) + if compute_ref: + policy.get_reference_policy_logprobs_from_meta( + meta, + timer=timer, + ) + + # Driver pulls only the per-token columns it needs + # for masking / advantage. Bulk (input_ids, multimodal, + # output_ids, attention_mask, position_ids) stays in + # TQ ā€” workers will fetch it via ``train_presharded``. + extras_bdd = policy.read_from_dataplane( + meta, + select_fields=[ + "prev_logprobs", + "generation_logprobs", + "token_mask", + *(["reference_policy_logprobs"] if compute_ref else []), + ], + pad_value_dict=_pad_dict, + ) + prev_logprobs = extras_bdd["prev_logprobs"] + generation_logprobs = extras_bdd["generation_logprobs"] + token_mask = extras_bdd["token_mask"] + reference_policy_logprobs = ( + extras_bdd["reference_policy_logprobs"] if compute_ref else None + ) + + # Thin BDD for the data-driven masking call: take + # the slice you need, transform, write delta back. + masking_data = BatchedDataDict[ClippedPGLossDataDict]( + { + "token_mask": token_mask, + "sample_mask": loss_multiplier, + "prev_logprobs": prev_logprobs, + "generation_logprobs": generation_logprobs, + } + ) + + ( + max_seq_mult_prob_error, + num_masked_seqs, + masked_correct_pct, + ) = compute_and_apply_seq_logprob_error_masking( + train_data=masking_data, + rewards=rewards, + seq_logprob_error_threshold=master_config.grpo[ + "seq_logprob_error_threshold" + ], + ) + # masking may have mutated sample_mask in place ā€” + # capture the post-masking value for delta-write. + sample_mask = masking_data["sample_mask"] + + with timer.time("advantage_calculation"): + print("ā–¶ Computing advantages...", flush=True) + mask = token_mask * sample_mask.unsqueeze(-1) + + # GRPO / Reinforce++ ignore ``repeated_batch`` (it's + # swallowed via ``**kwargs``); GDPO reads the + # per-component reward keys returned by + # ``get_gdpo_reward_component_keys``. The actor stashes + # those keys into ``driver_carry`` ā€” same payload as + # legacy passing the full repeated_batch. + adv_inputs = BatchedDataDict( + { + "total_reward": rewards, + "baseline": baseline, + "std": std, + } + ) + for k in get_gdpo_reward_component_keys(driver_carry): + adv_inputs[k] = driver_carry[k] + advantages = adv_estimator.compute_advantage( + prompt_ids=prompt_ids_for_adv, + rewards=rewards, + mask=mask, + repeated_batch=adv_inputs, + logprobs_policy=prev_logprobs, + logprobs_reference=reference_policy_logprobs, + ) + del prompt_ids_for_adv + + _log_mixed_rewards_and_advantages_information( + logger=logger, + total_steps=total_steps, + metrics=metrics, + baseline=baseline_for_log, + advantages=advantages, + ) + del baseline_for_log + + # ā”€ā”€ Driver delta-write: advantages + (post-masking) + # sample_mask under the same meta.sample_ids so workers fetch + # the union via train_presharded. + policy.write_to_dataplane( + meta, + fields={ + "advantages": advantages, + "sample_mask": sample_mask, + }, + ) + + memory_tracker.snapshot_start_of_stage("Policy train", dir()) + print("ā–¶ Preparing for training...", flush=True) + with timer.time("training_prep"): + policy.prepare_for_training() + POLICY_GENERATION_STALE = True + + print("ā–¶ Training policy...", flush=True) + with timer.time("policy_training"): + # Meta-driven train: workers fetch the union of + # rollout + driver-written + worker-written columns + # from TQ, train, return aggregated metrics via Ray. + train_results = policy.train_from_meta( + meta, + loss_fn=loss_fn, + timer=timer, + ) + + if sync_kv_scales: + with timer.time("recompute_kv_scales"): + print( + "ā–¶ Recomputing KV cache scales after policy update...", + flush=True, + ) + # Positive include-list ā€” calibration only consumes + # seq-dim tensor inputs. Train-side deltas + # (logprobs/advantages/masks) and wire-only message + # log bulk fields are skipped by virtue of not being + # in DP_CALIB_INPUT_FIELDS. + _calib_fields = [ + f for f in (meta.fields or []) if f in DP_CALIB_INPUT_FIELDS + ] + calibration_data = policy.read_from_dataplane( + meta, + select_fields=_calib_fields, + pad_value_dict=_pad_dict, + ) + kv_scales_cache = policy.calibrate_qkv_fp8_scales( + calibration_data, + include_q=True, + )["layers"] + POLICY_GENERATION_STALE = True + + # Stash input_ids and content before clear_samples so the + # late log_data jsonl block can use them. The clear below + # removes meta.sample_ids from TQ, so any post-clear + # read_columns on this meta would fail. ``content`` is a + # decoded object array (list[str]); read_columns decodes + # the NonTensorStack wire field via materialize. + _log_input_ids: Optional[torch.Tensor] = None + _log_content: Optional[np.ndarray] = None + if not _should_log_nemo_gym_responses(master_config): + _log_select = ["input_ids"] + if "content" in (meta.fields or []): + _log_select.append("content") + _log_extras = policy.read_from_dataplane( + meta, + select_fields=_log_select, + pad_value_dict=_pad_dict, + ) + _log_input_ids = _log_extras["input_ids"] + _log_content = _log_extras.get("content") + + # ā”€ā”€ Step-end TQ cleanup ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + policy.finish_step(meta) + + is_last_step = total_steps + 1 >= max_num_steps + if not master_config.data["use_multiple_dataloader"]: + is_last_step = is_last_step or ( + (current_epoch + 1 == max_num_epochs) + and (current_step + 1 == len(wrapped_dataloader)) + ) + + if (val_period > 0 and (total_steps + 1) % val_period == 0) or ( + val_at_end and is_last_step + ): + memory_tracker.snapshot_start_of_stage("Validation", dir()) + if NEED_REFIT and POLICY_GENERATION_STALE: + refit_policy_generation( + policy, + policy_generation, + colocated_inference, + kv_scales=kv_scales_cache if sync_kv_scales else None, + ) + POLICY_GENERATION_STALE = False + else: + if colocated_inference: + policy.offload_after_refit() + policy_generation.prepare_for_generation() + val_metrics, validation_timings = validate_sync( + rollout_actor=rollout_actor, + policy=policy, + val_dataloader=val_dataloader, + val_task_to_env=val_task_to_env, + step=total_steps + 1, + master_config=master_config, + logger=logger, + ) + policy_generation.finish_generation() + logger.log_metrics( + validation_timings, total_steps + 1, prefix="timing/validation" + ) + logger.log_metrics( + val_metrics, total_steps + 1, prefix="validation" + ) + + # advantages and token_mask are in scope from the + # advantage / masking blocks above. No need to re-fetch. + response_advantages = torch.masked_select(advantages, token_mask.bool()) + + memory_tracker.snapshot_start_of_stage("Metrics", dir()) + metrics = { + **metrics, + "loss": train_results["loss"].numpy(), + "grad_norm": train_results["grad_norm"].numpy(), + "reward": rewards.numpy(), + "mean_prompt_length": length.numpy(), + "total_num_tokens": input_lengths.numpy(), + "advantages/mean": torch.mean(response_advantages).detach().item() + if response_advantages.numel() > 0 + else 0.0, + "advantages/max": torch.max(response_advantages).detach().item() + if response_advantages.numel() > 0 + else 0.0, + "advantages/min": torch.min(response_advantages).detach().item() + if response_advantages.numel() > 0 + else 0.0, + **ds_metrics, + } + if "moe_metrics" in train_results: + metrics.update( + {f"moe/{k}": v for k, v in train_results["moe_metrics"].items()} + ) + # Cumulative unfiltered total_reward across all DS iterations + # (sliced to train_prompts_size). Falls back to filtered + # rewards if apply_dynamic_sampling didn't provide it + # (mid-step path). Hoisted once for reuse in metrics, jsonl, + # and the per-step print below. + unfiltered_rewards = ( + unfiltered_rewards_for_logging + if unfiltered_rewards_for_logging is not None + else rewards + ) + if master_config.grpo["use_dynamic_sampling"]: + metrics["filtered_reward"] = rewards.numpy() + metrics["reward"] = unfiltered_rewards.numpy() + + metrics.update(train_results["all_mb_metrics"]) + metrics.update(gen_step_metrics) + for k, v in metrics.items(): + if k in {"probs_ratio_min", "probs_ratio_clamped_min"}: + valid_values = [x for x in v if not np.isinf(x)] + metrics[k] = ( + np.min(valid_values).item() if valid_values else -1.0 + ) + elif k in {"probs_ratio_max", "probs_ratio_clamped_max"}: + valid_values = [x for x in v if not np.isinf(x)] + metrics[k] = ( + np.max(valid_values).item() if valid_values else -1.0 + ) + elif k in { + "lr", + "wd", + "reward", + "filtered_reward", + "global_valid_seqs", + "global_valid_toks", + "mean_prompt_length", + }: + metrics[k] = np.mean(v).item() + elif isinstance(v, (np.ndarray, list)): + metrics[k] = np.sum(v).item() + else: + print(f"Skipping aggregation for {k} ({type(v)})") + + metrics.update(rollout_metrics) + metrics["generation_logger_metrics"] = generation_logger_metrics + total_valid_tokens += metrics["global_valid_toks"] + + metrics["max_seq_mult_prob_error"] = max_seq_mult_prob_error + metrics["num_masked_seqs_by_logprob_error"] = num_masked_seqs + metrics["masked_correct_pct"] = masked_correct_pct + + consumed_samples += master_config.grpo["num_prompts_per_step"] + timeout.mark_iteration() + + should_save_by_step = ( + is_last_step + or (total_steps + 1) % master_config.checkpointing["save_period"] + == 0 + ) + should_save_by_timeout = timeout.check_save() + + memory_tracker.snapshot_start_of_stage("Checkpointing", dir()) + if master_config.checkpointing["enabled"] and ( + should_save_by_step or should_save_by_timeout + ): + policy.prepare_for_training() + + grpo_save_state["current_step"] = current_step + 1 + grpo_save_state["total_steps"] = total_steps + 1 + grpo_save_state["current_epoch"] = current_epoch + grpo_save_state["total_valid_tokens"] = total_valid_tokens + if val_metrics is not None: + grpo_save_state["val_reward"] = val_metrics["accuracy"] + elif "val_reward" in grpo_save_state: + del grpo_save_state["val_reward"] + grpo_save_state["consumed_samples"] = consumed_samples + + full_metric_name = master_config.checkpointing["metric_name"] + if full_metric_name is not None: + assert full_metric_name.startswith( + "train:" + ) or full_metric_name.startswith("val:"), ( + f"metric_name={full_metric_name} must start with 'val:' or 'train:'" + ) + prefix, metric_name = full_metric_name.split(":", 1) + metrics_source = metrics if prefix == "train" else val_metrics + if not metrics_source: + warnings.warn( + f"You asked to save checkpoints based on {metric_name} but no {prefix} metrics were collected. ", + stacklevel=2, + ) + if full_metric_name in grpo_save_state: + del grpo_save_state[full_metric_name] + elif metric_name not in metrics_source: + raise ValueError( + f"Metric {metric_name} not found in {prefix} metrics" + ) + else: + grpo_save_state[full_metric_name] = metrics_source[ + metric_name + ] + + with timer.time("checkpointing"): + print( + f"Saving checkpoint for step {total_steps + 1}...", + flush=True, + ) + checkpoint_path = checkpointer.init_tmp_checkpoint( + total_steps + 1, grpo_save_state, master_config + ) + policy.save_checkpoint( + weights_path=os.path.join( + checkpoint_path, "policy", "weights" + ), + optimizer_path=os.path.join( + checkpoint_path, "policy", "optimizer" + ) + if checkpointer.save_optimizer + else None, + tokenizer_path=os.path.join( + checkpoint_path, "policy", "tokenizer" + ), + checkpointing_cfg=master_config.checkpointing, + ) + if master_config.data["use_multiple_dataloader"]: + for ( + task_name, + task_dataloader, + ) in wrapped_dataloader.dataloaders.items(): + torch.save( + task_dataloader.state_dict(), + os.path.join( + checkpoint_path, + f"train_dataloader_{task_name}.pt", + ), + ) + else: + torch.save( + wrapped_dataloader.state_dict(), + os.path.join(checkpoint_path, "train_dataloader.pt"), + ) + checkpointer.finalize_checkpoint(checkpoint_path) + + memory_tracker.snapshot_start_of_stage("Logging", dir()) + # Per-step log_data jsonl. The 1-hop driver holds per-token + # slices it computed against (advantages, sample_mask, + # prev_logprobs, generation_logprobs, token_mask). For + # ``token_ids`` we fetch the small ``input_ids`` column from + # TQ at log time ā€” same data-driven slice pattern as masking + # / KV calibration. + if not _should_log_nemo_gym_responses(master_config): + log_data: dict = {} + if "agent_ref" in repeated_batch: + log_data["agent_ref"] = repeated_batch["agent_ref"] + if master_config.grpo["use_dynamic_sampling"]: + # Legacy semantics: ``rewards`` is unfiltered total_reward, + # ``filtered_rewards`` is the kept slice that's trained on. + log_data["rewards"] = unfiltered_rewards.tolist() + log_data["filtered_rewards"] = rewards.tolist() + else: + log_data["rewards"] = rewards.tolist() + log_data["input_lengths"] = input_lengths.tolist() + log_data["token_loss_mask"] = token_mask.tolist() + log_data["sample_loss_mask"] = sample_mask.tolist() + log_data["advantages"] = advantages.tolist() + log_data["generation_logprobs"] = generation_logprobs.tolist() + log_data["prev_logprobs"] = prev_logprobs.tolist() + # input_ids was stashed before the step-end clear_samples (the + # keys are no longer in TQ at this point); ``_log_input_ids`` + # is None when nemo_gym-responses logging path skipped the + # outer ``if not _should_log_nemo_gym_responses`` branch. + if _log_input_ids is not None: + log_data["token_ids"] = _log_input_ids.tolist() + # ``content`` (raw assistant text) is fetched from TQ as + # an object-array column above (stashed before clear_samples). + if _log_content is not None: + log_data["content"] = _log_content.tolist() + logger.log_batched_dict_as_jsonl( + log_data, f"train_data_step{total_steps + 1}.jsonl" + ) + del log_data + + timing_metrics: dict = timer.get_timing_metrics(reduction_op="sum") # type: ignore + if metrics["token_mult_prob_error"] > 1.05: + logger.log_plot_token_mult_prob_error( + { + "prompt_lengths": length, + "full_lengths": input_lengths, + "generation_logprobs": generation_logprobs, + "prev_logprobs": prev_logprobs, + "token_mask": token_mask, + "sample_mask": sample_mask, + }, + total_steps + 1, + name="train/token_mult_prob_error_plot_sample", + ) + if master_config.policy["generation"].get("vllm_cfg", {}).get( + "enable_vllm_metrics_logger", False + ) and master_config.logger.get("wandb_enabled", False): + log_generation_metrics_to_wandb( + generation_logger_metrics, + total_steps + 1, + master_config.policy["generation"]["vllm_cfg"][ + "vllm_metrics_logger_interval" + ], + logger, + ) + + if ( + master_config.policy["generation"] + .get("vllm_cfg", {}) + .get("async_engine", False) + ): + for metric_name in metrics.keys(): + if metric_name.startswith("histogram/"): + logger.log_histogram( + metrics[metric_name], + total_steps + 1, + f"generation_metrics/{metric_name}", + ) + + print("\nšŸ“Š Training Results:") + print(f" ā€¢ Loss: {metrics['loss']:.4f}") + if "draft_loss" in metrics: + print(f" ā€¢ Draft Loss: {metrics['draft_loss']:.4f}") + print(f" ā€¢ Generation KL Error: {metrics['gen_kl_error']:.4f}") + if master_config.grpo["use_dynamic_sampling"]: + print(f" ā€¢ Avg Filtered Reward: {np.mean(rewards.numpy()):.4f}") + print( + f" ā€¢ Avg Total Reward: {np.mean(unfiltered_rewards.numpy()):.4f}" + ) + else: + print(f" ā€¢ Avg Reward: {np.mean(rewards.numpy()):.4f}") + print( + f" ā€¢ Mean Generation Length: {metrics_logging_data['mean_gen_tokens_per_sample']:.4f}", + flush=True, + ) + + print("\nā±ļø Timing:", flush=True) + total_time = timing_metrics.get("total_step_time", 0) + + number_of_samples_per_step = ( + master_config.grpo["num_prompts_per_step"] + * master_config.grpo["num_generations_per_prompt"] + ) + total_num_gpus = ( + master_config.cluster["num_nodes"] + * master_config.cluster["gpus_per_node"] + ) + + print(f" ā€¢ Total step time: {total_time:.2f}s", flush=True) + + for k, v in sorted( + timing_metrics.items(), key=lambda item: item[1], reverse=True + ): + if k != "total_step_time": + percent = (v / total_time * 100) if total_time > 0 else 0 + print(f" ā€¢ {k}: {v:.2f}s ({percent:.1f}%)", flush=True) + + timing_metrics["valid_tokens_per_sec_per_gpu"] = ( + metrics["global_valid_toks"] / total_time / total_num_gpus + ) + performance_metrics = print_performance_metrics( + train_results, metrics, timing_metrics, master_config + ) + + logger.log_metrics(metrics, total_steps + 1, prefix="train") + logger.log_metrics( + performance_metrics, total_steps + 1, prefix="performance" + ) + logger.log_metrics( + timing_metrics, + total_steps + 1, + prefix="timing/train", + step_finished=True, + ) + + dynamic_sampling_num_gen_batches = 0 + + memory_tracker.snapshot_start_of_stage("After CPU memory clear", dir()) + + del repeated_batch + del rewards + del metrics + if "val_metrics" in dir(): + del val_metrics + + timer.reset() + current_step += 1 + total_steps += 1 + if should_save_by_timeout: + memory_tracker.snapshot_start_of_stage("", dir()) + print("Timeout has been reached, stopping training early", flush=True) + return + if total_steps >= max_num_steps: + memory_tracker.snapshot_start_of_stage("", dir()) + print( + "Max number of steps has been reached, stopping training early", + flush=True, + ) + return + + current_epoch += 1 + current_step = 0 diff --git a/nemo_rl/algorithms/reward_functions.py b/nemo_rl/algorithms/reward_functions.py index 974bebc6392..24547ecabe8 100644 --- a/nemo_rl/algorithms/reward_functions.py +++ b/nemo_rl/algorithms/reward_functions.py @@ -135,22 +135,34 @@ def apply_reward_shaping( # Calculate the expected response length expected_response_length = max_response_length - overlong_buffer_length - assert len(batch["message_log"]) == len(rewards), ( + # Prefer slim per-sample tensor (data-plane path: message_log lives in + # TQ, slice carries response_token_lengths). Fall back to scanning + # message_log for the legacy non-data-plane caller. + response_token_lengths = batch.get("response_token_lengths") + if response_token_lengths is not None: + if isinstance(response_token_lengths, torch.Tensor): + response_lengths = response_token_lengths.tolist() + else: + response_lengths = list(response_token_lengths) + else: + response_lengths = [] + for message_log in batch["message_log"]: + length = None + for message in message_log: + if message["role"] == "assistant": + length = message["token_ids"].shape[0] + break + assert length is not None, ( + "Assistant response not found during reward shaping" + ) + response_lengths.append(length) + + assert len(response_lengths) == len(rewards), ( "The number of messages in the batch must match the number of rewards" ) updated_rewards = torch.zeros_like(rewards) - for i, message_log in enumerate(batch["message_log"]): - # Get the assistant response length (index 1 is the assistant response) - message_response_length = None - for message in message_log: - if message["role"] == "assistant": - message_response_length = message["token_ids"].shape[0] - break - assert message_response_length is not None, ( - "Assistant response not found during reward shaping" - ) - + for i, message_response_length in enumerate(response_lengths): # Calculate the exceed length and the corresponding reward penalty exceed_length = message_response_length - expected_response_length overlong_reward = min( diff --git a/nemo_rl/data/llm_message_utils.py b/nemo_rl/data/llm_message_utils.py index 32bac1e9230..29840b2b8d2 100644 --- a/nemo_rl/data/llm_message_utils.py +++ b/nemo_rl/data/llm_message_utils.py @@ -14,6 +14,7 @@ import warnings from typing import Any, Optional, Union, cast +import numpy as np import torch from datasets import Dataset from transformers.tokenization_utils_base import PreTrainedTokenizerBase @@ -687,3 +688,135 @@ def remap_dataset_keys( lambda x: {v: x[k] for k, v in mapping_dict.items()}, remove_columns=list(mapping_dict.keys()), ) + + +# ā”€ā”€ Decomposed wire format for `message_log` ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ +# +# `message_log` mixes torch.Tensor with Python objects at the per-row +# level (`{"role": str, "content": str, "token_ids": Tensor, ...}` per +# turn). Shipping that shape per-row through pickle serializes the +# *underlying storage* of view-aliased tensor slices ā€” for a vllm batched +# output arena that's ~100 MB per row instead of the slice's ~10 KB. +# +# The helpers below split `message_log` into per-field arrays at the +# wire boundary (token tensors flat in `bulk_batch`, role/content +# strings as object arrays, per-turn lengths as one slim tensor) and +# rebuild the list-of-dicts shape on the consumer from local-arena +# views. No tensor ever reaches per-row pickle. + +# Fields ridden by `bulk_batch` and consumed by +# :func:`reconstruct_message_log` to rebuild the list-of-dicts view. +MESSAGE_LOG_BULK_FIELDS = ("turn_lengths", "turn_roles", "turn_contents") + + +def decompose_message_log( + message_log_batch: list[LLMMessageLogType], +) -> dict[str, Any]: + """Split a list-of-lists-of-dicts ``message_log`` into per-field arrays. + + Returns a dict with: + + - ``turn_lengths`` ā€” ``torch.LongTensor(B, max_turns)``, zero in unused slots. + - ``turn_roles`` ā€” ``np.ndarray(object, (B,))`` of ``list[str]``. + - ``turn_contents`` ā€” ``np.ndarray(object, (B,))`` of ``list[str]``. + - ``response_token_lengths`` ā€” ``torch.LongTensor(B,)``, assistant-turn + length per sample (0 if no assistant turn). Consumed by + :func:`nemo_rl.algorithms.reward_functions.apply_reward_shaping`. + """ + batch_size = len(message_log_batch) + max_turns = max((len(ml) for ml in message_log_batch), default=0) + + turn_roles = np.empty(batch_size, dtype=object) + turn_contents = np.empty(batch_size, dtype=object) + # Build Python lists in the hot loop; one tensor allocation at the end + # avoids per-turn 0-d tensor writes inside the loop. + turn_lengths_lol: list[list[int]] = [[0] * max_turns for _ in range(batch_size)] + response_lengths: list[int] = [0] * batch_size + + for i, ml in enumerate(message_log_batch): + roles: list[str] = [] + contents: list[str] = [] + lengths_i = turn_lengths_lol[i] + for t, m in enumerate(ml): + role = m["role"] # required; surface bad data loudly here + roles.append(role) + contents.append(m.get("content", "")) + tok = m.get("token_ids") + if tok is None: + continue + length = int(tok.shape[0]) if isinstance(tok, torch.Tensor) else len(tok) + lengths_i[t] = length + if role == "assistant" and response_lengths[i] == 0: + response_lengths[i] = length + turn_roles[i] = roles + turn_contents[i] = contents + + return { + "turn_lengths": torch.tensor(turn_lengths_lol, dtype=torch.long), + "turn_roles": turn_roles, + "turn_contents": turn_contents, + "response_token_lengths": torch.tensor(response_lengths, dtype=torch.long), + } + + +def attach_message_log_view(batch: BatchedDataDict[Any]) -> None: + """Attach ``batch['message_log']`` in place if decomposed fields are present. + + Rebuilds ``message_log`` as views into the consumer-local ``input_ids`` + / ``generation_logprobs``. Aliasing is harmless because the local + tensors own their storage and consumers do not re-pickle ``message_log``. + No-op when the decomposed fields are absent (legacy pickle-shipped path). + """ + if "input_ids" not in batch or any(k not in batch for k in MESSAGE_LOG_BULK_FIELDS): + return + batch["message_log"] = reconstruct_message_log( + input_ids=batch["input_ids"], + turn_lengths=batch["turn_lengths"], + turn_roles=batch["turn_roles"], + turn_contents=batch["turn_contents"], + generation_logprobs=batch.get("generation_logprobs"), + ) + + +def reconstruct_message_log( + input_ids: Tensor, + turn_lengths: Tensor, + turn_roles: "np.ndarray", + turn_contents: "np.ndarray", + generation_logprobs: Optional[Tensor] = None, +) -> list[LLMMessageLogType]: + """Inverse of :func:`decompose_message_log`. + + Per-turn ``token_ids`` and ``generation_logprobs`` are **views** into + the consumer-local ``input_ids`` / ``generation_logprobs`` tensors. + The aliasing is harmless because the local tensors own their storage + (decoded from the wire) and consumers do not re-pickle ``message_log``. + """ + batch_size = int(input_ids.shape[0]) + # Single host-side materialization ā€” avoids a per-turn .item() sync. + turn_lengths_list = turn_lengths.tolist() + out: list[LLMMessageLogType] = [] + for i in range(batch_size): + roles_i = turn_roles[i] + contents_i = turn_contents[i] + lengths_i = turn_lengths_list[i] + turns: LLMMessageLogType = [] + offset = 0 + for t, role in enumerate(roles_i): + length = lengths_i[t] + if length == 0: + turns.append({"role": role, "content": contents_i[t]}) + continue + turn: dict[str, Any] = { + "role": role, + "content": contents_i[t], + "token_ids": input_ids[i, offset : offset + length], + } + if generation_logprobs is not None and role == "assistant": + turn["generation_logprobs"] = generation_logprobs[ + i, offset : offset + length + ] + offset += length + turns.append(turn) + out.append(turns) + return out diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md new file mode 100644 index 00000000000..046b4e059d0 --- /dev/null +++ b/nemo_rl/data_plane/README.md @@ -0,0 +1,493 @@ +# nemo_rl.data_plane + +Stable boundary between NeMo-RL and the underlying data-plane backend +(currently `transfer_queue`; future: `nv-dataplane`). Every call site in +`nemo_rl/algorithms`, `nemo_rl/experience`, `nemo_rl/models` goes through +`DataPlaneClient`. No code imports `transfer_queue` directly outside the +adapter. + +--- + +## Vocabulary + +- **partition** ā€” a named data-flow scope in TQ (e.g. `"train"`, + `"val"`). Each partition owns its own field schema, consumer task + set, and per-sample production-status matrix. Sync GRPO uses one + stable partition (`"train"`) that is cleared and reused across + steps ā€” different partitions are for different data flows + (training vs validation vs replay buffer), not for different steps. +- **sample** ā€” one row in a partition, identified by a per-sample **key** + (e.g. `"_g0"`). Lives in TQ until `kv_clear`. +- **field** ā€” a named column (e.g. `input_ids`, `advantages`). Producers + write fields; consumers select them on read. Each `(sample, field)` + pair has an independent "produced?" bit on the TQ controller. +- **task** ā€” a *consumer* name (e.g. `"prev_lp"`, `"train"`). Each task + has its own consumption cursor, used by the task-mediated API only. +- **`KVBatchMeta`** ā€” the receipt returned by writes. Carries the keys, + partition id, sequence lengths, and the **fields written in this + put**. NOT a partition-wide schema view ā€” see the cheat-sheet below. + +--- + +## Mental model + +**TQ is a distributed storage and transfer engine.** It holds bulk +tensors (input_ids, logprobs, masks) addressed by per-sample keys, +moves them between producer and consumer Ray actors over the wire, +and tracks per-`(sample, field)` production status so consumers know +when their inputs are ready. Storage is transient: data lives in TQ +for the duration of one GRPO step and `kv_clear` drops it at step +end. The driver never holds bulk between rollout and training ā€” only +small per-sample slices (rewards, advantages) and metadata +(`KVBatchMeta`) cross the driver. + +**Three layers, one-way dependency:** + +``` +algorithms/grpo_sync.py ā† orchestration (sync trainer) + ā”‚ + ā–¼ +data_plane/{column_io, preshard} ā† producer/consumer helpers + ā”‚ + ā–¼ +data_plane/interfaces.py ā† stable boundary (DataPlaneClient) + ā”‚ + ā–¼ +data_plane/adapters/ ā† TransferQueue / NoOp / future nv-dataplane +``` + +--- + +## Legacy vs TQ-mediated ā€” same algorithm, encapsulated I/O + +The TQ-mediated trainer (`grpo_train_sync`) is meant to read like the +legacy in-memory trainer (`grpo_train`). The algorithm is identical; +only the data-fetch and lifecycle calls move behind `TQPolicy` / `meta` +methods. Per-step side-by-side: + +| Step | Legacy (`grpo.py: grpo_train`) | TQ-mediated (`grpo_sync.py: grpo_train_sync`) | +|---|---|---| +| Step start | (implicit) | `policy.prepare_step(N, group_size)` | +| Rollout | `run_multi_turn_rollout(...)` driver-side | `ray.get(rollout_actor.rollout_to_tq.remote(...))` ā€” bulk written to TQ inside the actor | +| Carry per-row data | `repeated_batch[k]` | `driver_carry[k]` (returned alongside `meta`) | +| Reward scale / shape / baseline / std | unchanged | unchanged | +| Mirror std for filter | `std` tensor in scope | `meta.stamp_tags({"std": ā€¦, "baseline": ā€¦})` | +| Dynamic sampling filter | `repeated_batch.select_indices(keep_idx)` | `meta.subset(keep_idx)` + `driver_carry.select_indices(keep_idx)` (inside `_apply_dynamic_sampling`, which also `kv_clear`s dropped uids) | +| Overlong filter / mask | unchanged | unchanged | +| Read columns for masking | `repeated_batch["generation_logprobs"]`, `repeated_batch["token_mask"]` | `policy.read_from_dataplane(meta, select_fields=["generation_logprobs", "token_mask"])` | +| Compute advantage | unchanged | unchanged | +| Write back advantage | mutate `repeated_batch["advantages"]` | `policy.write_to_dataplane(meta, {"advantages": ā€¦})` | +| Train | `policy.train(repeated_batch, loss_fn)` | `policy.train_from_meta(meta, loss_fn)` | +| Step end | (Python GC) | `policy.finish_step(meta)` | + +**The shape of the algorithm is unchanged.** Each TQ-mediated step has +a one-to-one counterpart in legacy; the only difference is where data +lives (Python memory vs TQ) and which method moves it. + +Per-stage audit grade after the encapsulation refactor: **A**. The +trainer body never references `policy.dp_client` directly ā€” only meta +and policy methods. `_apply_dynamic_sampling` still takes a raw +`dp_client` argument by design so unit tests can inject +`NoOpDataPlaneClient`. + +--- + +## E2E flow ā€” one sync GRPO step + +``` +ā”Œā”€ DRIVER Ā· grpo_train_sync ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā” +ā”‚ ā‘  policy.prepare_step(num_samples, group_size) ā”‚ +ā”‚ ā†’ register "train" partition with DP_TRAIN_FIELDS schema ā”‚ +ā”‚ ā‘” meta, driver_carry, *_ = ray.get( ā”‚ +ā”‚ rollout_actor.rollout_to_tq.remote(repeated_batch, uids=ā€¦)) ā”‚ +ā”‚ ā† single Ray RPC; actor runs rollout + flatten + mask + ā”‚ +ā”‚ kv_first_write of bulk under uid-derived keys. ā”‚ +ā””ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”¬ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”˜ + ā”‚ bulk now in TQ; driver has meta + driver_carry slice + ā–¼ +ā”Œā”€ DRIVER (reward + advantage, on driver_carry only) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā” +ā”‚ ā‘¢ scale_rewards / apply_reward_shaping (legacy parity) ā”‚ +ā”‚ ā‘£ baseline, std = calculate_baseline_and_std_per_prompt(...) ā”‚ +ā”‚ meta.stamp_tags({"std": ā€¦, "baseline": ā€¦}) ā”‚ +ā”‚ ā†’ filter-without-fetch primitive on meta ā”‚ +ā”‚ ā‘¤ [optional] _apply_dynamic_sampling(meta, driver_carry, ā€¦) ā”‚ +ā”‚ ā†’ meta.subset(keep) + driver_carry.select_indices(keep) ā”‚ +ā”‚ ā†’ dp_client.kv_clear(dropped_keys) ā”‚ +ā”‚ ā‘„ overlong filter (loss_multiplier = 0 on truncated rows) ā”‚ +ā””ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”¬ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”˜ + ā–¼ +ā”Œā”€ DRIVER ā†’ WORKERS (logprob phase) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā” +ā”‚ ā‘¦ prev_lp = policy.get_logprobs_from_meta(meta) ā”‚ +ā”‚ ref_lp = policy.get_reference_policy_logprobs_from_meta(meta) ā”‚ +ā”‚ ā†“ inside the policy method: ā”‚ +ā”‚ shard_meta_for_dp(meta) ā€” length-balanced split, pure meta ā”‚ +ā”‚ fan-out: worker.get_logprobs_presharded.remote(shard) Ɨ N ā”‚ +ā”‚ ā†’ _fetch(shard) ā†’ kv_batch_get ā†’ materialize ā”‚ +ā”‚ ā†’ forward ā†’ logprobs ā”‚ +ā”‚ ā†’ leader writes back as new TQ column on meta.keys ā”‚ +ā”‚ ā‘§ extras = policy.read_from_dataplane(meta, select_fields=[ā€¦]) ā”‚ +ā”‚ advantages = compute_advantages(...) ā”‚ +ā”‚ ā‘Ø policy.write_to_dataplane(meta, {"advantages": ā€¦, "sample_mask":ā€¦})ā”‚ +ā””ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”¬ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”˜ + ā–¼ +ā”Œā”€ DRIVER ā†’ WORKERS (train + cleanup) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā” +ā”‚ ā‘© policy.train_from_meta(meta, loss_fn=ā€¦) ā”‚ +ā”‚ ā†“ same shard_meta_for_dp + fan-out shape; no write-back ā”‚ +ā”‚ (training is terminal). ā”‚ +ā”‚ ā‘Ŗ policy.finish_step(meta) ā†’ drop step's bulk from TQ ā”‚ +ā””ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”˜ + ā†’ next step ā†’ ā‘  +``` + +Bulk tensors live in TQ; the driver only holds `meta` + the small +`driver_carry` slice. On-wire layout is jagged +(`codec.pack_jagged_fields` ā†” `codec.materialize` at every put / get). + +--- + +## `KVBatchMeta` + +The receipt for a put. `meta.fields` is only what was written by *this* +put, not the partition-wide schema. See `interfaces.py` for the ABC. + +| Attribute | Meaning | +|---|---| +| `partition_id` | TQ partition these keys live in | +| `keys` | Per-sample row identifiers | +| `fields` | Fields written by the put that minted this meta | +| `sequence_lengths` | Per-row valid (unpadded) lengths ā€” drives length-balanced sharding | +| `tags` | `list[dict]` 1:1 with `keys` ā€” per-row primitive sidecar for filter-without-fetch | +| `extra_info` | Batch-level bag (`rollout_metrics`, `pad_to_multiple`, `global_forward_pad_seqlen`, packing metadata) | +| `task_name` | Optional consumer tag, carried through | + +**Hard rules** ā€” `kv_batch_put` fields must be `TensorDict` of tensors +(or `np.ndarray(dtype=object)`); primitives go on `tags`. `select_fields` +is required on every `kv_batch_get` ā€” no implicit "fetch all". + +--- + +## Helpers above the client + +| Helper | What it does | +|---|---| +| `column_io.kv_first_write` | Rollout actor's flat first put. Caller mints `keys`. | +| `column_io.read_columns` / `write_columns` | `kv_batch_get` / `kv_batch_put` + jagged ā†” padded materialize. | +| `preshard.shard_meta_for_dp` | Pure metadata split, length-balanced when packing args are passed. | +| `KVBatchMeta.subset` / `.slice` / `.concat` | Pure meta transforms used by dynamic sampling; thread `tags` 1:1 with `keys`. | +| `KVBatchMeta.stamp_tags` | Mirror per-row scalars onto `meta.tags`. Init-if-None + length check. | +| `codec.pack_jagged_fields` | Jagged-pack at every put boundary. | + +--- + +## Per-sample key invariant + +Keys are minted **once** at rollout (`key_i = f"{uid}_g{i}"`) and reused +for every subsequent `kv_batch_put` / `kv_batch_get` on that sample. +Worker write-backs append new columns under the same keys. + +--- + +## Concrete examples + +### Call shapes + +A real step at production scale ā€” +`num_prompts_per_step=128, num_generations_per_prompt=4`, DP world = 8, +prompt ā‰ˆ 512 tok, response ā‰¤ 1024 tok. Final batch is `128 Ɨ 4 = 512` +rows. + +**1. Step prepare + rollout** (driver ā€” `grpo_train_sync` body): + +```python +# Open the per-step TQ partition. Cleared and reused across steps. +policy.prepare_step(num_samples=512, group_size=4) + +# One Ray RPC bundles: clear gen metrics ā†’ rollout ā†’ flatten + mask ā†’ +# kv_first_write of bulk to TQ ā†’ finish_generation ā†’ metrics snapshot. +# The actor handles 6 stages internally; the driver gets back the +# meta handle + a small per-row tensor slice. +n_prompts = repeated_batch.size # 512 (= 128 prompts Ɨ 4 gens) +uids = [str(uuid.uuid4()) for _ in range(n_prompts // 4)] # 128 uids +meta, driver_carry, rollout_metrics, gen_metrics = ray.get( + rollout_actor.rollout_to_tq.remote( + repeated_batch, + uids=uids, + partition_id=policy.tq_partition_id, # "train" + first_iter=(dynamic_sampling_num_gen_batches == 1), + ) +) +# meta.keys ā‰ˆ ["a3f9_g0", "a3f9_g1", "a3f9_g2", "a3f9_g3", +# "b7c1_g0", ā€¦] (512 keys) +# meta.sequence_lengths ā‰ˆ [847, 612, 1503, 989, 711, ā€¦] (actual lens) +# meta.fields = ["input_ids", "input_lengths", +# "generation_logprobs", "token_mask", +# "sample_mask", ā€¦multimodal extrasā€¦] +# driver_carry : BatchedDataDict of per-row tensors +# (total_reward, loss_multiplier, truncated, +# length, input_lengths, prompt_ids_for_adv, +# response_token_lengths, GDPO components) +``` + +**2. Reward + dynamic sampling** (driver, on `driver_carry` only): + +```python +driver_carry = scale_rewards(driver_carry, cfg["grpo"]["reward_scaling"]) +if cfg["grpo"]["reward_shaping"]["enabled"]: + driver_carry = apply_reward_shaping(driver_carry, cfg["grpo"]["reward_shaping"]) +driver_carry["baseline"], driver_carry["std"] = ( + calculate_baseline_and_std_per_prompt( + driver_carry["prompt_ids_for_adv"], + driver_carry["total_reward"], + torch.ones_like(driver_carry["total_reward"]), + leave_one_out_baseline=cfg["grpo"]["use_leave_one_out_baseline"], + ) +) +# Mirror std/baseline onto meta so dynamic sampling can filter on +# meta alone (no tensor fetch). +meta.stamp_tags( + { + "std": driver_carry["std"].tolist(), + "baseline": driver_carry["baseline"].tolist(), + } +) + +# DAPO non-zero-std filter ā€” drops rows where the prompt's reward +# variance is zero, kv_clears their bulk, accumulates survivors +# across iterations until train_prompts_size (512) is reached. +if cfg["grpo"]["use_dynamic_sampling"]: + pending_meta, pending_carry, *_ = _apply_dynamic_sampling( + meta=meta, driver_carry=driver_carry, + pending_meta=pending_meta, pending_carry=pending_carry, + train_prompts_size=512, + num_gen_batches=dynamic_sampling_num_gen_batches, + max_gen_batches=cfg["grpo"]["dynamic_sampling_max_gen_batches"], + dp_client=policy.dp_client, + ) +``` + +**3. Logprob + advantage + write-back**: + +```python +# Worker fan-out happens inside these. Per-DP-rank shard via +# shard_meta_for_dp(meta, dp_world=8, ā€¦); each worker fetches its +# ~64 keys via kv_batch_get and writes back the result column under +# the same keys on the leader. +prev_lp = policy.get_logprobs_from_meta(meta, timer=timer)["logprobs"] +ref_lp = policy.get_reference_policy_logprobs_from_meta(meta, timer=timer) +ref_lp = ref_lp["reference_logprobs"] + +# Driver-side per-token columns for masking. Tiny delta ā€” just two +# fields Ɨ 512 rows. +extras = policy.read_from_dataplane( + meta, + select_fields=["generation_logprobs", "token_mask"], + pad_value_dict=_pad_dict, +) +advantages = adv_estimator.compute_advantage( + prompt_ids=driver_carry["prompt_ids_for_adv"], + rewards=rewards, mask=mask, + repeated_batch=adv_inputs, + logprobs_policy=prev_lp, + logprobs_reference=ref_lp, +) + +# Write the per-token advantage + post-masking sample_mask back to TQ +# under meta.keys so workers fetch the unified view in train. +policy.write_to_dataplane( + meta, + fields={"advantages": advantages, "sample_mask": sample_mask}, +) +``` + +**4. Train + cleanup**: + +```python +train_results = policy.train_from_meta(meta, loss_fn=loss_fn, timer=timer) +policy.finish_step(meta) # drop step's bulk from TQ +``` + +**5. Validation path** ā€” slim `driver_carry` to skip ~1 MB/batch: + +```python +# inside validate_sync; val_batch_size ā‰ˆ 64 +policy.prepare_val_partition(n_prompts, partition_id="val") +meta, driver_carry, rollout_metrics, _ = ray.get( + rollout_actor.rollout_to_tq.remote( + val_batch, uids=uids, partition_id="val", + finish_generation=False, # keep inference state warm + task_to_env_override=val_task_to_env, + carry_keys=["total_reward"], # only field val consumes + ) +) +total_rewards.extend(driver_carry["total_reward"].tolist()) +mlog_cols = policy.read_from_dataplane( + meta, select_fields=["turn_roles", "turn_contents"], +) +policy.finish_step(meta) +``` + +### Sequence-length flow (seqpack / dynbatch) + +How `meta.sequence_lengths` routes samples to DP ranks. Worked example +sized to one production microbatch ā€” 4 prompts Ɨ 2 generations = 8 +samples, DP world = 4, lengths typical of math/code rollouts. + +``` +# Rollout actor flattens prompt + response per sample. +# input_lengths[i] = prompt_len_i + response_len_i (actual content, +# unpadded). +sample 0 (a3f9_g0): prompt=312, response= 892 ā†’ input_lengths=1204 +sample 1 (a3f9_g1): prompt=312, response= 187 ā†’ input_lengths= 499 +sample 2 (b7c1_g0): prompt=421, response= 1024 ā†’ input_lengths=1445 ā† long +sample 3 (b7c1_g1): prompt=421, response= 455 ā†’ input_lengths= 876 +sample 4 (c0d8_g0): prompt=148, response= 213 ā†’ input_lengths= 361 ā† short +sample 5 (c0d8_g1): prompt=148, response= 339 ā†’ input_lengths= 487 +sample 6 (d2e1_g0): prompt=276, response= 651 ā†’ input_lengths= 927 +sample 7 (d2e1_g1): prompt=276, response= 402 ā†’ input_lengths= 678 + +# kv_first_write returns meta row-aligned with keys: +meta.keys = ["a3f9_g0", "a3f9_g1", "b7c1_g0", "b7c1_g1", + "c0d8_g0", "c0d8_g1", "d2e1_g0", "d2e1_g1"] +meta.sequence_lengths = [ 1204, 499, 1445, 876, + 361, 487, 927, 678 ] + +# shard_meta_for_dp slices keys + sequence_lengths with the SAME +# idx_list ā€” driver-side, no TQ I/O. Length-balanced via seqpack: +rank 0: idx=[2, 4] ā†’ keys=["b7c1_g0","c0d8_g0"] lens=[1445, 361] = 1806 +rank 1: idx=[0, 5] ā†’ keys=["a3f9_g0","c0d8_g1"] lens=[1204, 487] = 1691 +rank 2: idx=[6, 1] ā†’ keys=["d2e1_g0","a3f9_g1"] lens=[ 927, 499] = 1426 +rank 3: idx=[3, 7] ā†’ keys=["b7c1_g1","d2e1_g1"] lens=[ 876, 678] = 1554 +# Ī£ packed lengths per rank within ~25% ā€” well-balanced. + +# Each worker fetches its own ~64 keys per step from TQ: +data = self._fetch(shard) # kv_batch_get(shard.keys, select_fields=ā€¦) +``` + +**Gotcha ā€” `make_sequence_length_divisible_by` (TPƗCP alignment)**: +`input_ids` is padded to a multiple of TPƗCP at write time (e.g. 8 for +TP=4, CP=2), but `input_lengths` is the actual content length. Seqpack +balances on actual lengths; padding is reapplied per shard. + +``` +# row with input_lengths=1204, TPƗCP=8 ā†’ input_ids padded to 1208: +input_ids: [t0, t1, ā€¦, t1203, 0, 0, 0, 0] # 1208 elems +input_lengths: 1204 # actual +meta.sequence_lengths: 1204 # what seqpack uses āœ“ +``` + +**Gotcha ā€” DP-rank seq-dim alignment (`global_forward_pad_seqlen`)**: +Each DP rank's `_fetch` would otherwise pad to its slice's local max, +so two ranks in the same step could forward at different seq dims. +That breaks any collective that assumes cross-rank shape uniformity +(mcore MoE all-to-all, CP, etc.). The data plane handles this with a +single per-batch cap minted on the driver: + +* `TQPolicy._stamp_pad_seqlen(meta)` runs before every fan-out + (`train_from_meta`, `_logprob_dispatch`, `read_from_dataplane`). + Idempotent ā€” sets `meta.extra_info["global_forward_pad_seqlen"]` + to `round_up(max(meta.sequence_lengths), max(pad_to_multiple, + sequence_length_round))` on first call, no-op on subsequent calls. +* `shard_meta_for_dp` propagates `extra_info` to every per-rank meta + via `dict(meta.extra_info)` ā€” so all ranks see the same target. +* Worker `_fetch` and driver `read_columns` both pass + `pad_to_seqlen = meta.extra_info["global_forward_pad_seqlen"]` + into `codec.materialize`, which right-pads the seq dim to that + absolute target. All DP ranks within a step therefore return + columns at one identical seq dim. + +Opt out in tests with `_fetch(..., dp_aligned_seq_len=False)` to +observe per-rank local-pad behavior. + +``` +# 4 DP ranks, slice maxes: [1208, 1320, 944, 1080]; sequence_length_round=64 +global_forward_pad_seqlen = round_up(1320, 64) = 1344 +# All 4 ranks pad their materialized tensors to seq_dim=1344. +``` + +--- + +## Configuration + +The data plane is configured via a `data_plane:` block in the master +YAML (`examples/configs/...`). **YAML is the single source of truth +for defaults** ā€” the adapter has no hidden `cfg.get(key, default)` +fallbacks. The canonical exemplar is +`examples/configs/grpo_math_1B.yaml`. + +All eight keys below are **required** when `enabled=true`. Recipes +under `examples/configs/recipes/**/*.yaml` inherit them via +`defaults:` from the exemplar. + +```yaml +data_plane: + enabled: false # flip to true to engage grpo_train_sync + impl: transfer_queue # only one impl today + backend: "simple" # "simple" or "mooncake_cpu" + storage_capacity: 1000000 # max samples retained per partition + num_storage_units: 2 # storage shards + claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence + global_segment_size: 549755813888 # 512 GiB ā€” used when backend == "mooncake_cpu" + local_buffer_size: 68719476736 # 64 GiB ā€” used when backend == "mooncake_cpu" + # observability: # NotRequired + # enabled: false +``` + +Backend choice: +- **`simple`** ā€” ZMQ-backed; lowest setup overhead. Default for tests + and small runs. +- **`mooncake_cpu`** ā€” Mooncake transfer engine; higher throughput at + scale. Required for multi-node clusters with large bulk volume. + +Capacity rule of thumb (any backend): + +``` +storage_capacity ā‰„ 2 Ɨ num_prompts Ɨ n_gens Ɨ max_seq_len + Ɨ bytes_per_token Ɨ num_active_fields +``` + +The `2 Ɨ` headroom covers dynamic sampling overflow and one step of +pipelining between rollout and training. + +--- + +## When `data_plane.enabled=False` + +`build_data_plane_client` raises ā€” there is no NoOp prod fallback. +For the no-data-plane path use the legacy +`nemo_rl.algorithms.grpo.grpo_train`; the sync trainer +`grpo_train_sync` requires `enabled=True` and a `TQPolicy`. + +`NoOpDataPlaneClient` (`adapters/noop.py`) exists only as a unit-test +fixture for the ABC contract tests. + +--- + +## Where to look + +| Concern | File | +|---|---| +| Stable boundary (ABC) | `nemo_rl/data_plane/interfaces.py` | +| Adapter (TransferQueue impl) | `nemo_rl/data_plane/adapters/transfer_queue.py` | +| Adapter (NoOp, test only) | `nemo_rl/data_plane/adapters/noop.py` | +| Codec (jagged pack / unpack) | `nemo_rl/data_plane/codec.py` | +| Column-level helpers | `nemo_rl/data_plane/column_io.py` (`read_columns`, `write_columns`, `kv_first_write`) | +| DP-rank meta sharding | `nemo_rl/data_plane/preshard.py` | +| Worker fetch + leader write-back | `nemo_rl/data_plane/worker_mixin.py` | +| Schema constants | `nemo_rl/data_plane/schema.py` | +| Rollout actor (first put) | `nemo_rl/experience/sync_rollout_actor.py` | +| TQ-mediated Policy subclass | `nemo_rl/models/policy/tq_policy.py` | +| End-to-end orchestration | `nemo_rl/algorithms/grpo_sync.py` | +| Unit tests | `tests/data_plane/unit/` | +| Functional tests (real backends) | `tests/data_plane/functional/` | + +--- + +## Async path (proposed) + +The data-plane interface covers both sync and async, but the **sync +trainer uses only half of it**. The task-mediated half +(`claim_meta` / `get_data` / `check_consumption_status`) is reserved +for the async trainer, which is not yet wired into production. + +Design proposal, filtering / staleness strategies, and open questions: +see [`docs/data-plane-async-proposal.md`](docs/data-plane-async-proposal.md). diff --git a/nemo_rl/data_plane/__init__.py b/nemo_rl/data_plane/__init__.py new file mode 100644 index 00000000000..56b19178a1c --- /dev/null +++ b/nemo_rl/data_plane/__init__.py @@ -0,0 +1,38 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""NeMo-RL data-plane package. + +The public surface is intentionally tiny: an ABC, a meta dataclass, a +config TypedDict, and a factory. Everything else is an implementation +detail of a specific adapter. +""" + +from nemo_rl.data_plane.codec import materialize +from nemo_rl.data_plane.factory import build_data_plane_client +from nemo_rl.data_plane.interfaces import ( + DataPlaneClient, + DataPlaneConfig, + KVBatchMeta, +) +from nemo_rl.data_plane.observability import MetricsDataPlaneClient, log_event + +__all__ = [ + "DataPlaneClient", + "DataPlaneConfig", + "KVBatchMeta", + "MetricsDataPlaneClient", + "build_data_plane_client", + "log_event", + "materialize", +] diff --git a/nemo_rl/data_plane/adapters/__init__.py b/nemo_rl/data_plane/adapters/__init__.py new file mode 100644 index 00000000000..341a77c5bc6 --- /dev/null +++ b/nemo_rl/data_plane/adapters/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/nemo_rl/data_plane/adapters/noop.py b/nemo_rl/data_plane/adapters/noop.py new file mode 100644 index 00000000000..1c5b00a5e44 --- /dev/null +++ b/nemo_rl/data_plane/adapters/noop.py @@ -0,0 +1,244 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""In-memory ``DataPlaneClient`` test fixture. + +Behaves like a real adapter end-to-end (put ā†’ get ā†’ clear, consumption +counters, field-presence as the stage-done signal) but stores everything +in process memory. The ABC contract tests run against this implementation +so they don't require TQ installed. + +Production callers must NOT use this ā€” :func:`build_data_plane_client` +intentionally raises when ``enabled=False`` rather than returning a NoOp +fallback (see ``factory.py``). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch +from tensordict import TensorDict + +from nemo_rl.data_plane.codec import stack_or_nest as _stack_or_nest +from nemo_rl.data_plane.interfaces import DataPlaneClient, KVBatchMeta + + +def _reject_non_tensor_leaves(td: TensorDict) -> None: + """No pickle on the bus. Mirror of the TQ adapter check. + + Walk the leaves via ``keys()`` + indexed lookup rather than + ``items()``, because some tensordict versions skip ``NonTensorData`` + entries from ``items(leaves_only=True)`` ā€” they're "leaves" by + structure but not tensor-typed, so they'd silently slip past a + naive items() iteration. + """ + bad = [] + for k in td.keys(include_nested=True, leaves_only=True): + v = td.get(k) + if not isinstance(v, torch.Tensor): + bad.append(k) + if bad: + raise TypeError( + f"put_samples received non-tensor leaves: {bad}. " + "Tensorize via codec helpers, use `tags=` for primitives, " + "or use the Ray object store for arbitrary Python objects." + ) + + +@dataclass +class _Partition: + fields: list[str] + num_samples: int + consumer_tasks: list[str] + grpo_group_size: int | None + enums: dict[str, list[str]] + rows: dict[str, dict[str, torch.Tensor]] = field(default_factory=dict) + tags: dict[str, dict[str, Any]] = field(default_factory=dict) + # per-task set of keys already returned by claim_meta (TQ ``mode='fetch'``) + consumed: dict[str, set[str]] = field(default_factory=dict) + + +class NoOpDataPlaneClient(DataPlaneClient): + """Reference in-memory implementation.""" + + def __init__(self) -> None: + self._partitions: dict[str, _Partition] = {} + self._closed = False + + def register_partition( + self, + partition_id: str, + fields: list[str], + num_samples: int, + consumer_tasks: list[str], + grpo_group_size: int | None = None, + enums: dict[str, list[str]] | None = None, + ) -> None: + self._partitions[partition_id] = _Partition( + fields=list(fields), + num_samples=int(num_samples), + consumer_tasks=list(consumer_tasks), + grpo_group_size=grpo_group_size, + enums=dict(enums) if enums else {}, + consumed={t: set() for t in consumer_tasks}, + ) + + def claim_meta( + self, + partition_id: str, + task_name: str, + required_fields: list[str], + batch_size: int, + dp_rank: int | None = None, + blocking: bool = True, + timeout_s: float = 60.0, + ) -> KVBatchMeta: + del blocking, timeout_s, dp_rank # NoOp is single-process + rec = self._partitions[partition_id] + if task_name not in rec.consumed: + raise KeyError( + f"task {task_name!r} not registered as a consumer of " + f"partition {partition_id!r}" + ) + + ready: list[str] = [] + seqs: list[int] = [] + for key, row in rec.rows.items(): + if key in rec.consumed[task_name]: + continue + if not all(f in row for f in required_fields): + continue + ready.append(key) + tag = rec.tags.get(key, {}) + seqs.append(int(tag.get("input_lengths", 0))) + if len(ready) >= batch_size: + break + + rec.consumed[task_name].update(ready) + return KVBatchMeta( + partition_id=partition_id, + task_name=task_name, + sample_ids=ready, + fields=list(required_fields), + sequence_lengths=seqs if any(seqs) else None, + ) + + def get_data( + self, + meta: KVBatchMeta, + select_fields: list[str] | None = None, + ) -> TensorDict: + fields = select_fields if select_fields is not None else meta.fields + if fields is None: + raise ValueError( + "get_data requires either select_fields or meta.fields; " + "fetching all fields silently is forbidden." + ) + return self.get_samples(meta.sample_ids, meta.partition_id, list(fields)) + + def check_consumption_status( + self, partition_id: str, task_names: list[str] + ) -> bool: + rec = self._partitions[partition_id] + for t in task_names: + if t not in rec.consumed: + return False + if len(rec.consumed[t]) < len(rec.rows): + return False + return True + + def put_samples( + self, + sample_ids: list[str], + partition_id: str, + fields: TensorDict | None = None, + tags: list[dict[str, Any]] | None = None, + ) -> KVBatchMeta: + rec = self._partitions[partition_id] + if fields is not None: + _reject_non_tensor_leaves(fields) + for i, sid in enumerate(sample_ids): + row = rec.rows.setdefault(sid, {}) + for fname in fields.keys(): + val = fields[fname][i] + # Defense in depth ā€” _reject_non_tensor_leaves can + # miss NonTensorData entries depending on the + # tensordict version's iteration semantics. + if not isinstance(val, torch.Tensor): + raise TypeError( + f"put_samples received non-tensor leaf " + f"{fname!r}: {type(val).__name__}. " + "Tensorize via codec helpers, use `tags=` " + "for primitives, or use the Ray object store " + "for arbitrary Python objects." + ) + row[fname] = val.detach().clone() + if tags is not None: + for sid, tag in zip(sample_ids, tags): + rec.tags.setdefault(sid, {}).update(tag) + return KVBatchMeta( + partition_id=partition_id, + task_name=None, + sample_ids=list(sample_ids), + fields=list(fields.keys()) if fields is not None else None, + tags=[dict(t) for t in tags] if tags is not None else None, + ) + + def get_samples( + self, + sample_ids: list[str], + partition_id: str, + select_fields: list[str], + ) -> TensorDict: + rec = self._partitions[partition_id] + if not sample_ids: + return TensorDict({}, batch_size=(0,)) + + out: dict[str, list[torch.Tensor]] = {f: [] for f in select_fields} + for sid in sample_ids: + row = rec.rows[sid] + for f in select_fields: + if f not in row: + raise KeyError( + f"field {f!r} not yet produced for sample_id {sid!r} " + f"in partition {partition_id!r}" + ) + out[f].append(row[f]) + + stacked = {f: _stack_or_nest(out[f]) for f in select_fields} + return TensorDict(stacked, batch_size=(len(sample_ids),)) + + def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None: + rec = self._partitions.get(partition_id) + if rec is None: + return + if sample_ids is None: + rec.rows.clear() + rec.tags.clear() + for s in rec.consumed.values(): + s.clear() + self._partitions.pop(partition_id, None) + return + for sid in sample_ids: + rec.rows.pop(sid, None) + rec.tags.pop(sid, None) + for s in rec.consumed.values(): + s.discard(sid) + + def close(self) -> None: + if self._closed: + return + self._partitions.clear() + self._closed = True diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py new file mode 100644 index 00000000000..963c4917638 --- /dev/null +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -0,0 +1,660 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Adapter wiring :class:`DataPlaneClient` onto the ``transfer_queue`` package. + +Pure plumbing ā€” it owns the TQ controller / client handle and translates +:class:`KVBatchMeta` ā†” TQ's own ``BatchMeta`` / ``KVBatchMeta``. No +business logic. Backend init is lifted from +``rl-arena/arena/backends.py``; the call shapes are lifted from +``rl-arena/arena/dataplane_client.py``. +""" + +from __future__ import annotations + +import ipaddress +import os +import socket +import subprocess +import time +from importlib import resources +from typing import Any + +import torch +import transfer_queue as tq +from tensordict import TensorDict + +from nemo_rl.data_plane.interfaces import ( + DataPlaneClient, + DataPlaneConfig, + KVBatchMeta, +) + +# ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ +# Backend init ā€” lifted from rl-arena/arena/backends.py. +# ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def _get_local_node_ip() -> str: + """Return THIS process's host IP, not the cluster head's. + + Each Ray actor process must use its own node's IP so Mooncake's + announce address (``MC_TCP_BIND_ADDRESS`` ā†’ ``desc.ip_or_host_name`` + in ``transfer_engine_impl.cpp``) is routable cross-node. + Non-routable addresses are rejected: + + * Link-local (169.254/16, fe80::/10) ā€” ``gethostbyname`` can + resolve to APIPA on hosts where ``avahi-autoipd`` is active. + * Loopback (127.0.0.0/8, ::1) ā€” hosts whose ``/etc/hosts`` maps + the hostname to 127.0.0.1 would otherwise announce an + unroutable address to Mooncake peers, causing cross-node + ``connection refused``. + """ + try: + ip = socket.gethostbyname(socket.gethostname()) + addr = ipaddress.ip_address(ip) + if addr.is_link_local or addr.is_loopback: + return "" + return ip + except Exception: + return "" + + +def _mooncake_transport_config() -> dict: + protocol = os.environ.get("MC_MOONCAKE_PROTOCOL", "tcp") + if protocol != "rdma": + return {"protocol": "tcp"} + device = os.environ.get("MC_MOONCAKE_DEVICE", "") + if not device: + try: + out = subprocess.run( + [ + "sh", + "-c", + "for d in /sys/class/infiniband/mlx5_*/ports/1/link_layer; do " + " test -f $d && grep -q Ethernet $d && basename $(dirname $(dirname $d)); " + "done | head -1", + ], + check=False, + capture_output=True, + text=True, + ).stdout.strip() + device = out or "" + except Exception: + device = "" + if device: + os.environ.setdefault("MC_GID_INDEX", os.environ.get("MC_GID_INDEX", "3")) + return {"protocol": "rdma", "device_name": device} + + +def _connect_existing() -> None: + """Worker-process path: connect this process's client to the Ray cluster. + + Connects to the already-running named controller actor. Mirrors + rl-arena/arena/dataplane_client.py's `tq.init()` (no args) call. + """ + tq.init() + + +_TQ_RUNTIME_ENV_PATCHED = False + + +def _resolve_tq_pin() -> str: + """Return the ``TransferQueue`` requirement string from nemo-rl metadata. + + Single source of truth is ``pyproject.toml`` ā€” we read it back via + ``importlib.metadata.requires`` so the runtime_env injection cannot + drift from the dependency declaration. + """ + from importlib.metadata import requires + + for req in requires("nemo-rl") or []: + spec = req.split(";")[0].strip() + if spec.lower().startswith("transferqueue"): + return spec + raise RuntimeError( + "Could not resolve TransferQueue dependency from nemo-rl metadata. " + "Check pyproject.toml under [project.dependencies]." + ) + + +def _patch_tq_actor_runtime_env() -> None: + """Inject a per-actor ``runtime_env`` pin into TQ's actor ``.options()``. + + TQ spawns ``SimpleStorageUnit`` and ``TransferQueueController`` via + ``Cls.options(...).remote(...)`` without a runtime_env, so they + inherit the job-level env. In a multi-node container deployment + where each node has its own ``/opt/nemo_rl_venv``, the driver's + ``uv sync`` only updates ray-head's venv and a worker-node actor + fails with ``ModuleNotFoundError``. This monkey-patch makes Ray + pip-install TQ into a per-actor runtime_env on first spawn (cached + per-node by Ray afterwards). Idempotent. Couples us to TQ's internal + class layout ā€” if TQ restructures, this becomes a no-op with a + logged warning and we fall back to per-node ``uv sync``. + + The pin is sourced from nemo-rl's installed metadata via + :func:`_resolve_tq_pin` so it cannot drift from ``pyproject.toml``. + + TODO(zhiyul): remove this patch once the nightly container image + is published with ``TransferQueue`` baked in via ``pyproject.toml``. + When every node starts from that image, the base env already has TQ + and Ray actors inherit it ā€” this injection then becomes pure + overhead (Ray builds a redundant per-actor pip env on top of the + container's existing TQ install). Drop the call from + ``TQDataPlaneClient.__init__`` and delete this function. + """ + global _TQ_RUNTIME_ENV_PATCHED + if _TQ_RUNTIME_ENV_PATCHED: + return + + runtime_env = {"pip": [_resolve_tq_pin()]} + + def _install(cls) -> bool: + if not hasattr(cls, "options"): + return False + original = cls.options + + def patched(*args, **kwargs): + kwargs.setdefault("runtime_env", runtime_env) + return original(*args, **kwargs) + + cls.options = patched # type: ignore[method-assign] + return True + + patched_any = False + try: + from transfer_queue.storage.simple_backend import SimpleStorageUnit + + patched_any |= _install(SimpleStorageUnit) + except ImportError: + pass + try: + from transfer_queue.controller import TransferQueueController + + patched_any |= _install(TransferQueueController) + except ImportError: + pass + + if not patched_any: + # Soft-fail: TQ may have moved its actor classes. The driver will + # still work; multi-node TQ may need the per-node `uv sync` workaround. + import warnings + + warnings.warn( + "Could not patch TQ actor classes for runtime_env injection. " + "Multi-node TQ may fail with ModuleNotFoundError: 'transfer_queue' " + "on worker nodes. Workaround: run `uv sync` inside each node's " + "container before the driver runs.", + RuntimeWarning, + stacklevel=2, + ) + _TQ_RUNTIME_ENV_PATCHED = True + + +def _init_tq(cfg: DataPlaneConfig) -> None: + """Driver-process path: bootstrap the TQ controller for the chosen backend.""" + from omegaconf import OmegaConf + + base = OmegaConf.load(str(resources.files("transfer_queue") / "config.yaml")) + + backend = cfg["backend"] + storage_capacity = cfg["storage_capacity"] + num_storage_units = cfg["num_storage_units"] + + # polling_mode=True: controller returns empty BatchMeta instead of raising + # TimeoutError when no samples are ready yet. The client-side blocking + # loop in `claim_meta` drives the retry cadence. + controller_overlay = {"controller": {"polling_mode": True}} + + if backend == "simple": + overlay = { + **controller_overlay, + "backend": { + "storage_backend": "SimpleStorage", + "SimpleStorage": { + "total_storage_size": storage_capacity, + "num_data_storage_units": num_storage_units, + }, + }, + } + elif backend == "mooncake_cpu": + # The mooncake-transfer-engine wheel ships `mooncake_master` at + # /mooncake/, NOT on $PATH. TQ's + # subprocess.Popen(["mooncake_master", ...]) fails with + # FileNotFoundError unless we put the package dir on PATH first. + import mooncake # type: ignore[import-not-found] + + # TQ's mooncake_client masks any underlying ImportError as + # "Please install via pip install mooncake-transfer-engine". + # Force the real cause (e.g. ``libcudart.so.X: cannot open + # shared object file``) to surface by importing here. + import mooncake.store # type: ignore[import-not-found] # noqa: F401 + + _moon_pkg = os.path.dirname(mooncake.__file__) + _master = os.path.join(_moon_pkg, "mooncake_master") + try: + os.chmod(_master, 0o755) + except OSError as e: + if not os.access(_master, os.X_OK): + raise RuntimeError( + f"Failed to make {_master} executable: {e}. " + f"Mooncake bootstrap requires this binary." + ) from e + _existing_path = os.environ.get("PATH", "") + if _moon_pkg not in _existing_path.split(os.pathsep): + os.environ["PATH"] = _moon_pkg + os.pathsep + _existing_path + # Per-process MC_TCP_BIND_ADDRESS / KV-path promotion already + # set by TQDataPlaneClient.__init__ (runs on every process, + # including this driver). _init_tq only needs local_ip below + # for the metadata/master server URLs (driver-bound). + local_ip = _get_local_node_ip() + if not local_ip: + raise RuntimeError( + "Mooncake backend requires a local node IP; " + "_get_local_node_ip() returned empty." + ) + # Mooncake virtual segment / local buffer sizing. Defaults sized + # for production-scale rollouts (multi-iter DAPO, large + # message_log object payloads); under-sized values cause + # ``batch_get_tensor returned None`` once mooncake exhausts its + # internal allocator headroom. Lazy-mmap'd, so RSS is bounded + # by actual traffic. Override per-recipe via + # ``data_plane.global_segment_size`` / + # ``data_plane.local_buffer_size`` (bytes). + overlay = { + **controller_overlay, + "backend": { + "storage_backend": "MooncakeStore", + "MooncakeStore": { + "global_segment_size": int(cfg["global_segment_size"]), + "local_buffer_size": int(cfg["local_buffer_size"]), + # _init_tq runs on the driver only ā€” driver IS the + # head, so local_ip here is also the head's IP that + # mooncake_master + the metadata server bind to. + "metadata_server": f"{local_ip}:50050", + "master_server_address": f"{local_ip}:50051", + **_mooncake_transport_config(), + }, + }, + } + else: + raise ValueError(f"unknown TQ backend: {backend!r}") + + conf = OmegaConf.merge(base, overlay) + + # Inject runtime_env into TQ's actor spawn so SimpleStorageUnit / + # TransferQueueController land on workers with transfer_queue available + # ā€” see _patch_tq_actor_runtime_env() docstring for the why. + _patch_tq_actor_runtime_env() + + # pyrefly: ignore # bad-argument-type + tq.init(conf=conf) + + +# ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ +# Adapter-level enforcement that nothing but tensors crosses the bus. +# ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def _assert_no_key_loss(src_dict: dict, new_td: TensorDict, fn: str) -> None: + """Guard against silent leaf drops through TensorDict constructor rebuild. + + tensordict's constructor has historically dropped NonTensorStack / + NonTensorData leaves when built from a plain dict. Compare the + source dict's keys against the rebuilt TD's top-level keys. + """ + new_keys = set(new_td.keys()) + if set(src_dict.keys()) != new_keys: + dropped = sorted(set(src_dict.keys()) - new_keys) + raise RuntimeError( + f"{fn} lost leaves through TensorDict rebuild: dropped={dropped}." + ) + + +def _promote_1d_leaves(td: TensorDict) -> TensorDict: + """Unsqueeze 1D tensor leaves to ``(N, 1)`` ā€” mooncake_cpu KV-path workaround. + + Works around TQ's ``KVStorageManager`` 1D schema/data mismatch; + :func:`_from_wire` squeezes the trailing 1 back on read. Symmetric + with `_from_wire` ā€” callers gate on ``self._promote_1d``. + ``NonTensorStack`` / ``NonTensorData`` leaves pass through. + + Args: + td: ``TensorDict`` whose 1D tensor leaves should be promoted. + + Returns: + ``TensorDict`` with 1D tensor leaves unsqueezed to ``(N, 1)``; + all other leaves pass through unchanged. + """ + # td.keys() (top-level) includes NonTensorData / NonTensorStack leaves. + # keys(include_nested=True, leaves_only=True) enumerates tensor leaves + # only ā€” non-tensor leaves would silently fall out of the rebuilt dict. + new_dict: dict[str, Any] = {} + changed = False + for k in td.keys(): + v = td.get(k) + if isinstance(v, torch.Tensor) and not v.is_nested and v.dim() == 1: + new_dict[str(k)] = v.unsqueeze(-1).contiguous() + changed = True + else: + new_dict[str(k)] = v + if not changed: + return td + new_td = TensorDict(new_dict, batch_size=td.batch_size) + _assert_no_key_loss(new_dict, new_td, "_promote_1d_leaves") + return new_td + + +def _from_wire(td: TensorDict) -> TensorDict: + """Inverse of `_promote_1d_leaves`: squeeze trailing 1 back to (N,).""" + # Same top-level iteration as `_promote_1d_leaves`: NonTensorData / + # NonTensorStack leaves are only visible via td.keys(), not leaves_only. + new_dict: dict[str, Any] = {} + changed = False + for k in td.keys(): + v = td.get(k) + if ( + isinstance(v, torch.Tensor) + and not v.is_nested + and v.dim() >= 2 + and v.shape[-1] == 1 + ): + new_dict[str(k)] = v.squeeze(-1).contiguous() + changed = True + else: + new_dict[str(k)] = v + if not changed: + return td + new_td = TensorDict(new_dict, batch_size=td.batch_size) + _assert_no_key_loss(new_dict, new_td, "_from_wire") + return new_td + + +class TQDataPlaneClient(DataPlaneClient): + """Adapter faƧade ā€” maps NeMo-RL calls onto TransferQueue's public API.""" + + def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None: + """Construct a TQ-backed client. + + Args: + cfg: data-plane config (backend selection, poll cadence, ā€¦). + bootstrap: True (driver) bootstraps the TQ controller using + ``cfg``. False (worker) connects this process to an + already-running named controller actor in the Ray + cluster ā€” ``cfg`` is then only consulted for client-side + knobs (poll interval). + """ + # mooncake_cpu setup must run BEFORE _init_tq / _connect_existing + # ā€” once tq.init/connect runs, Mooncake's engine.so reads the + # env vars and they can't be changed. Three per-process knobs + # needed in EVERY process that builds a TQ client (driver, + # SyncRolloutActor, every MegatronPolicyWorker rank): + # 1. MC_TCP_BIND_ADDRESS ā€” Mooncake engine.so writes this into + # desc.ip_or_host_name, the address peers receive from the + # metadata service. Without it, getifaddrs()[0] picks usb0 + # (169.254.x APIPA) and peers fail to connect. + # 2. MC_STORE_MEMCPY=0 ā€” Mooncake LOCAL_MEMCPY fast-path + # reinterpret_casts cross-process pointers, segfaulting + # MemcpyWorkerPool. PR #1995 (merged 2026-04-30) fixes the + # root cause but isn't in any published wheel yet + # (mooncake-transfer-engine 0.3.10.post2 was bumped before + # that merge). Drop this once the wheel includes the fix. + # 3. KV-path 1D promotion ā€” works around TQ's + # extract_field_schema schema/data mismatch for 1D fields. + if cfg["backend"] == "mooncake_cpu": + local_ip = _get_local_node_ip() + if local_ip: + # Force-assign per-process: Ray actors inherit env vars + # from the driver, so a setdefault on the worker would + # be a no-op and the actor would announce the driver's + # IP ā€” peers fail with "connection refused". + os.environ["MC_TCP_BIND_ADDRESS"] = local_ip + os.environ.setdefault("MC_STORE_MEMCPY", "0") + + # Workaround for TQ KVStorageManager's 1D-field schema/data + # mismatch (only `mooncake_cpu` goes through that path; `simple` + # is unaffected). Writer unsqueezes 1D ā†’ (N, 1) on put; reader + # squeezes the trailing 1 back on get. Drop when upstream TQ + # unifies the schema/data shapes for 1D fields. + self._promote_1d = cfg["backend"] == "mooncake_cpu" + + if bootstrap: + _init_tq(cfg) + else: + _connect_existing() + self._poll_interval_s = cfg["claim_meta_poll_interval_s"] + self._closed = False + + # ā”€ā”€ (A) task-mediated ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + def register_partition( + self, + partition_id: str, + fields: list[str], + num_samples: int, + consumer_tasks: list[str], + grpo_group_size: int | None = None, + enums: dict[str, list[str]] | None = None, + ) -> None: + # Pre-populate ``Partition.field_name_mapping`` with the full + # field schema by doing a single synchronous placeholder put on + # the driver before any worker producer/consumer is live for + # this partition. + # + # Why: TQ's controller registers new field names lazily inside + # ``update_production_status`` (controller.py:538) without a lock, + # while ``kv_retrieve_meta`` (controller.py:1645) iterates the + # same dict ā€” interleaved threads raise ``RuntimeError: dictionary + # changed size during iteration`` and kill the controller's + # ProcessRequestThread (no try/except around the while-loop). + # Registering everything from a single driver thread before any + # client request races with a put removes the trigger entirely. + if not fields: + return + client = tq.get_client() + dummy_td = TensorDict( + {f: torch.zeros(1) for f in fields}, + batch_size=[1], + ) + meta = client.put(data=dummy_td, partition_id=partition_id) + client.clear_samples(metadata=meta) + + def claim_meta( + self, + partition_id: str, + task_name: str, + required_fields: list[str], + batch_size: int, + dp_rank: int | None = None, + blocking: bool = True, + timeout_s: float = 60.0, + ) -> KVBatchMeta: + client = tq.get_client() + deadline = time.time() + max(0.0, timeout_s) + sampling_config: dict[str, Any] = {} + if dp_rank is not None: + sampling_config["dp_rank"] = dp_rank + + while True: + tq_meta = client.get_meta( + data_fields=list(required_fields), + batch_size=int(batch_size), + partition_id=partition_id, + task_name=task_name, + mode="fetch", + sampling_config=sampling_config, + ) + if getattr(tq_meta, "size", 0) > 0: + break + if not blocking: + return KVBatchMeta( + partition_id=partition_id, + task_name=task_name, + sample_ids=[], + fields=list(required_fields), + ) + if time.time() >= deadline: + raise TimeoutError( + f"claim_meta(partition={partition_id}, task={task_name}) " + f"timed out after {timeout_s}s" + ) + time.sleep(self._poll_interval_s) + + keys: list[str] = client.kv_retrieve_keys( + global_indexes=list(tq_meta.global_indexes), + partition_id=partition_id, + ) + + # Propagate per-key tags. ``sequence_lengths`` is lifted out of + # the ``input_lengths`` tag if present (kept as a typed list + # because shard_meta_for_dp reads it directly), but the rest + # of the tag dict travels through unchanged so consumers can + # filter on it without fetching data. + tags = list(tq_meta.custom_meta) if tq_meta.custom_meta else [{} for _ in keys] + seqlens: list[int] | None = None + if tags and any("input_lengths" in t for t in tags): + seqlens = [int(t.get("input_lengths", 0)) for t in tags] + + return KVBatchMeta( + partition_id=partition_id, + task_name=task_name, + sample_ids=keys, + fields=list(required_fields), + sequence_lengths=seqlens, + tags=tags if tags else None, + ) + + def get_data( + self, + meta: KVBatchMeta, + select_fields: list[str] | None = None, + ) -> TensorDict: + fields = select_fields if select_fields is not None else meta.fields + if fields is None: + raise ValueError( + "get_data requires either select_fields or meta.fields; " + "silently fetching all fields is forbidden." + ) + return self.get_samples(meta.sample_ids, meta.partition_id, list(fields)) + + def check_consumption_status( + self, partition_id: str, task_names: list[str] + ) -> bool: + client = tq.get_client() + for t in task_names: + if not client.check_consumption_status( + task_name=t, partition_id=partition_id + ): + return False + return True + + # ā”€ā”€ (B) direct-by-key ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + def put_samples( + self, + sample_ids: list[str], + partition_id: str, + fields: TensorDict | None = None, + tags: list[dict[str, Any]] | None = None, + ) -> KVBatchMeta: + if not sample_ids: + return KVBatchMeta( + partition_id=partition_id, task_name=None, sample_ids=[], fields=None + ) + if tags is None: + tags = [{} for _ in sample_ids] + + wire_fields: TensorDict | None = None + field_names: list[str] | None = None + if fields is not None: + # No ``.contiguous()``: under tensordict==0.12.2 it strips + # non-tensor leaves (NonTensorStack stored as LinkedList) to empty + # TDs. TQ's encoder forces ``.contiguous()`` per tensor leaf + # itself, so the call here was redundant for tensors and + # destructive for non-tensors. + wire_fields = fields.detach() # type: ignore[bad-assignment,missing-argument] + if self._promote_1d: + wire_fields = _promote_1d_leaves(wire_fields) # type: ignore[bad-argument-type] + field_names = list(wire_fields.keys()) + + # TQ's wire vocabulary is `keys=` ā€” translation point. + tq.kv_batch_put( + keys=list(sample_ids), + partition_id=partition_id, + fields=wire_fields, + tags=tags, + ) + + return KVBatchMeta( + partition_id=partition_id, + task_name=None, + sample_ids=list(sample_ids), + fields=field_names, + tags=[dict(t) for t in tags] if tags else None, + ) + + def get_samples( + self, + sample_ids: list[str], + partition_id: str, + select_fields: list[str], + ) -> TensorDict: + if not sample_ids: + return TensorDict({}, batch_size=(0,)) + # TQ's wire vocabulary is `keys=` ā€” translation point. + td = tq.kv_batch_get( + keys=list(sample_ids), + partition_id=partition_id, + select_fields=select_fields, + ) + if self._promote_1d: + td = _from_wire(td) + return td + + def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None: + cleared_via_none = sample_ids is None + if sample_ids is None: + # No local state ā€” ask TQ's controller for the current key + # set in this partition. ``kv_list`` errors propagate; we + # don't want a network blip to silently turn into "cleared + # nothing". + listing = tq.kv_list(partition_id=partition_id) + sample_ids = list(listing.get(partition_id, {}).keys()) + if not sample_ids: + if cleared_via_none: + import warnings + + warnings.warn( + f"clear_samples(sample_ids=None, partition_id={partition_id!r}) " + "found nothing to clear ā€” TQ's kv_list returned no keys for " + "this partition. The partition may already be empty, never " + "have been written to, or be unknown to the controller. " + "Callers that hold a ``KVBatchMeta`` should pass its " + "``sample_ids`` explicitly for a deterministic clear.", + RuntimeWarning, + stacklevel=2, + ) + return + # TQ's wire vocabulary is `keys=` ā€” translation point. + tq.kv_clear(keys=list(sample_ids), partition_id=partition_id) + + # ā”€ā”€ (C) lifecycle ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + def close(self) -> None: + if self._closed: + return + self._closed = True + try: + tq.close() + except Exception: + pass diff --git a/nemo_rl/data_plane/codec.py b/nemo_rl/data_plane/codec.py new file mode 100644 index 00000000000..fb625687256 --- /dev/null +++ b/nemo_rl/data_plane/codec.py @@ -0,0 +1,372 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Wire <-> trainer codec ā€” jagged-on-the-wire bridge. + +* Writer side: variable-length fields are encoded as +``torch.nested.nested_tensor`` with ``layout=torch.jagged`` before +``put_samples``. Padding tax is paid only when a consumer needs a +rectangular tensor. + +* Reader side: :func:`materialize` accepts the wire TensorDict and, +when ``layout='padded'``, calls +:func:`torch.nested.to_padded_tensor` on any nested leaves using +the per-field padding value supplied in ``pad_value_dict``. Trainer +code consumes the padded BatchedDataDict unchanged. + +* Worker write-backs that produce ``response``-shaped outputs use +:func:`response_from_nested` to extract the response slice from a +(prompt+response) nested tensor. + +* Non-tensor object fields ride as ``NonTensorStack`` / ``NonTensorData`` +leaves (TQ-native passthrough). :func:`materialize` decodes them back +to ``np.ndarray(dtype=object)`` for the trainer. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch +from tensordict import TensorDict, TensorDictBase + +from nemo_rl.data_plane.schema import Layout + +if TYPE_CHECKING: + # Type-only import. At runtime, BatchedDataDict is loaded lazily + # inside materialize() ā€” see comment there for rationale. + from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + +# ā”€ā”€ Padded ā†” nested helpers ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def to_nested_by_length( + padded: torch.Tensor, + lengths: torch.Tensor, +) -> torch.Tensor: + """Strip right-padding off a rectangular tensor using per-row lengths. + + Used by the producer side: convert + :func:`batched_message_log_to_flat_message` output (already padded) + into the wire format before ``put_samples``. + + Args: + padded: Rectangular tensor of shape ``(N, S, ...)``. + lengths: Per-row valid lengths, shape ``(N,)``. CUDA tensors are + moved to CPU once to avoid per-row syncs. + + Returns: + A ``torch.jagged`` nested tensor whose i-th row is + ``padded[i, :lengths[i], ...]``. + """ + if padded.dim() < 2: + raise ValueError( + f"to_nested_by_length expects (N, S, ...); got shape {tuple(padded.shape)}" + ) + n = padded.shape[0] + if lengths.shape != (n,): + raise ValueError( + f"lengths shape {tuple(lengths.shape)} != ({n},) (rows of padded)" + ) + # Single sync ā€” without this, the per-row ``.item()`` below would + # GPU-sync N times if ``lengths`` lives on CUDA. + lens = lengths.cpu().tolist() if lengths.is_cuda else lengths.tolist() + rows = [padded[i, : lens[i]] for i in range(n)] + return torch.nested.as_nested_tensor(rows, layout=torch.jagged) + + +def stack_or_nest(tensors: list[torch.Tensor]) -> torch.Tensor: + """Stack equal-shape rows; reconstruct as jagged nested when ragged. + + Args: + tensors: Per-row tensors; assumed to share leading dims modulo + an optional ragged seq dim. Empty list returns ``torch.empty(0)``. + + Returns: + A regular tensor when all rows share shape; otherwise a + ``torch.jagged`` nested tensor. + """ + if not tensors: + return torch.empty(0) + first_shape = tensors[0].shape + if all(t.shape == first_shape for t in tensors): + return torch.stack(tensors, dim=0) + return torch.nested.as_nested_tensor(tensors, layout=torch.jagged) + + +def unwrap_wire_stripped_payload(item: Any) -> Any: + """Recover the payload of a possibly wire-stripped ``NonTensorData``. + + TQ's ``MsgpackEncoder._encode_tensordict`` serializes any + ``TensorDictBase`` via ``dict(obj.items())`` ā€” only the tensor + backing dict. ``NonTensorData`` stores its payload in + ``_non_tensordict["data"]``, so it round-trips through ZMQ as an + empty ``TensorDict({}, batch_size=[])``. We map only that exact + signature to ``None``; any other ``TensorDictBase`` (with tensor + fields, non-scalar batch, or a salvageable ``_non_tensordict`` + payload) passes through unchanged so we never drop real data. + """ + nt = getattr(item, "_non_tensordict", None) + if isinstance(nt, dict) and "data" in nt: + return nt["data"] + if ( + isinstance(item, TensorDictBase) + and item.batch_dims == 0 + and len(item.keys()) == 0 + ): + return None + return item + + +def maybe_pack_jagged( + val: torch.Tensor, + lengths: torch.Tensor, +) -> torch.Tensor: + """Convert ``val`` to jagged iff it looks like a per-token field. + + Used by every write site (initial put, driver delta-write, worker + write-back) so all per-token fields land in TQ as jagged with the + same row lengths ā€” read-time materialization then pads them all to + the same target shape, avoiding shape-mismatch crashes between + mixed wire formats. + + Args: + val: Tensor to consider. Qualifies for jagged conversion only + when ``val.shape == (N, max(lengths), ...)`` where + ``N == lengths.shape[0]``. + lengths: Per-row valid lengths, shape ``(N,)``. + + Returns: + A ``torch.jagged`` nested tensor when the shape heuristic matches; + otherwise ``val`` passed through as a rectangular tensor. + """ + n = lengths.shape[0] + if n == 0: + return val.detach().contiguous() + max_len = int(lengths.max().item()) + if val.dim() < 2 or val.shape[0] != n or val.shape[1] != max_len: + return val.detach().contiguous() + return to_nested_by_length(val.detach(), lengths) + + +def pack_jagged_fields( + fields: "dict[str, torch.Tensor | np.ndarray]", + *, + lengths: torch.Tensor | None, +) -> TensorDict: + """Pack a column dict into the wire layout expected by ``put_samples``. + + Zero-copy where possible: per-token tensors that match + ``(N, max(lengths), ...)`` become ``torch.jagged`` views via + :func:`maybe_pack_jagged`; non-conforming tensors pass through + rectangular; ``np.ndarray(dtype=object)`` is forwarded as-is. This + is a **layout transform**, not serialization ā€” the on-wire bytes are + produced later by the TQ backend's msgpack encoder. Centralizing + the transform here makes it the single source of truth for both + :func:`kv_first_write` and :func:`write_columns`. + + Args: + fields: Column name ā†’ tensor or object array. Other value types + raise ``TypeError``. + lengths: Per-row valid lengths used by :func:`maybe_pack_jagged` + to decide whether a tensor qualifies for jagged conversion. + ``None`` disables jagged conversion entirely (every tensor + passes through rectangular). + + Returns: + ``TensorDict`` with ``batch_size=[N]`` (N from ``lengths`` if + given, else 0) ready for ``put_samples``. + """ + n = int(lengths.shape[0]) if lengths is not None else 0 + packed: dict[str, Any] = {} + for k, v in fields.items(): + if isinstance(v, np.ndarray) and v.dtype == object: + # tensordict==0.12.2 wire bug: a NonTensorStack stored as a + # TensorDict leaf returns as a LinkedList on parent + # __getitem__, losing identity. ndarray(dtype=object) + # round-trips intact. + packed[k] = v + elif isinstance(v, torch.Tensor): + packed[k] = ( + maybe_pack_jagged(v, lengths) + if lengths is not None + else v.detach().contiguous() + ) + else: + raise TypeError( + f"pack_jagged_fields: unsupported value type for {k!r}: {type(v)}. " + "Use torch.Tensor or np.ndarray(dtype=object)." + ) + return TensorDict(packed, batch_size=[n]) + + +def pack_per_token_field(val: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor: + """Force-jaggedize a known per-token field, tolerating SP padding. + + Unlike :func:`maybe_pack_jagged` (which is shape-strict to avoid + false positives on 3D extras like image features), this function is + invoked at write-back sites where the caller already knows the + field is per-token (e.g. ``prev_logprobs``, + ``reference_policy_logprobs``). mcore SP rounds the forward + output's seq dim up to a multiple of TP, so the value can be 1+ + tokens wider than ``max(lengths)``; :func:`to_nested_by_length` + slices each row to its own length and drops the trailing SP + padding cleanly. + + Args: + val: Per-token tensor. Falls back to rectangular when it cannot + be jaggedized (wrong batch dim, < 2D, or seq dim shorter + than ``max(lengths)``). + lengths: Per-row valid lengths, shape ``(N,)``. + + Returns: + A ``torch.jagged`` nested tensor when the shape allows; + otherwise ``val`` passed through as a rectangular tensor. + """ + n = lengths.shape[0] + if n == 0: + return val.detach().contiguous() + max_len = int(lengths.max().item()) + if val.dim() < 2 or val.shape[0] != n or val.shape[1] < max_len: + return val.detach().contiguous() + return to_nested_by_length(val.detach(), lengths) + + +def response_from_nested( + full: torch.Tensor, + response_mask: torch.Tensor, +) -> torch.Tensor: + """Extract the response slice from a (prompt+response) nested tensor. + + Used on the worker side for logprob / ref-logprob write-back where + only the response-token slice is interesting downstream. The + "left-shift by one token" convention is applied (so logprobs at + output position i correspond to the prediction of input token i+1). + + Args: + full: Jagged nested tensor of shape + ``(N, prompt_len + response_len)``. + response_mask: Jagged nested tensor of shape + ``(N, response_len)``; its ``offsets().diff()`` gives the + per-row response length. + + Returns: + Jagged nested tensor of shape ``(N, response_len)`` containing + the left-shifted response slice. + """ + values = full.values() + offsets = full.offsets() + response_lens = response_mask.offsets().diff() + response_list = [] + for resp_len, seq_offset in zip(response_lens, offsets[1:], strict=True): + # left-shift output by one token for log_probs / values + response_list.append(values[seq_offset - resp_len - 1 : seq_offset - 1]) + return torch.nested.as_nested_tensor(response_list, layout=torch.jagged) + + +# ā”€ā”€ materialize: wire TensorDict ā†’ trainer BatchedDataDict ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def materialize( + td: TensorDict, + layout: Layout = "padded", + pad_value_dict: dict[str, int | float] | None = None, + pad_to_seqlen: int = 0, +) -> "BatchedDataDict[Any]": + """Convert a wire TensorDict to a BatchedDataDict. + + Trainer/worker code expects rectangular tensors ā€” this is the + bridge from the on-wire nested format. + + The lazy ``BatchedDataDict`` import keeps + ``import nemo_rl.data_plane`` cheap for unit tests that don't + actually call this function (``BatchedDataDict`` transitively + pulls multimodal deps like decord / torchvision). + + Args: + td: Wire TensorDict to materialize. + layout: ``"padded"`` (default) pads nested-tensor leaves via + :func:`torch.nested.to_padded_tensor` using + ``pad_value_dict[k]`` (or 0 if unspecified); rectangular + leaves pass through. ``"jagged"`` passes nested leaves + through ā€” use only when the caller knows how to consume + them. + pad_value_dict: Per-field pad value used when ``layout='padded'``. + pad_to_seqlen: When > 0, right-pad the seq dim up to this + absolute length after ``to_padded_tensor``. Worker-side + ``_fetch`` passes its forward-pass target here (rounded up + to ``sequence_length_round`` for Megatron's microbatch + iterator); driver-side ``read_columns`` leaves it 0 and + consumes the natural-padded shape. Default 0 disables. + + Returns: + ``BatchedDataDict`` with rectangular tensors for padded layout, + nested tensors for jagged layout, and ``np.ndarray(dtype=object)`` + for ``NonTensorStack`` leaves (TQ-native non-tensor passthrough). + """ + from tensordict import NonTensorData, NonTensorStack + + from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + pads = pad_value_dict or {} + out: dict[str, Any] = {} + # pyrefly: inference cycle on tensordict.items() loop var. + for key, val in td.items(include_nested=False): # type: ignore[bad-assignment] + if isinstance(val, NonTensorStack): + # ``np.asarray(list, dtype=object)`` would probe each item's + # ``__iter__`` to detect a nested array. A wire-stripped TD + # has ``batch_dims=0`` ā†’ its ``__iter__`` raises + # ``StopIteration`` ā†’ ``RuntimeError: generator raised + # StopIteration``. ``np.empty + assignment`` skips that + # probe; ``unwrap_wire_stripped_payload`` normalizes both + # live ``NonTensorData`` and stripped TDs. + items = val.tolist() + arr = np.empty(len(items), dtype=object) + for i, item in enumerate(items): + arr[i] = unwrap_wire_stripped_payload(item) + out[key] = arr + continue + if isinstance(val, NonTensorData): + out[key] = np.asarray([val.data], dtype=object) + continue + if not isinstance(val, torch.Tensor): + raise TypeError( + f"materialize() received unexpected leaf type for {key!r}: " + f"{type(val)}. Expected Tensor or NonTensorStack." + ) + if val.is_nested and layout == "padded": + pad = pads.get(key, 0) + padded = torch.nested.to_padded_tensor(val, padding=pad) + else: + pad = pads.get(key, 0) + padded = val + # Apply `pad_to_seqlen` to ALL 2D+ tensors, not only the freshly- + # padded-from-nested case. Rectangular wire payloads (vLLM's + # right-padded output) ride the ``else`` branch above, so without + # this they'd skip the cross-DP forward pad target and break the + # microbatch iterator (truncate_tensors ā†’ narrow length>size). + if ( + pad_to_seqlen > 0 + and isinstance(padded, torch.Tensor) + and padded.dim() >= 2 + and padded.shape[1] < pad_to_seqlen + ): + pad_spec = [0, 0] * (padded.dim() - 2) + [ + 0, + pad_to_seqlen - padded.shape[1], + ] + padded = torch.nn.functional.pad(padded, pad_spec, value=pad) + out[key] = padded + return BatchedDataDict(out) diff --git a/nemo_rl/data_plane/column_io.py b/nemo_rl/data_plane/column_io.py new file mode 100644 index 00000000000..9d2df9b9909 --- /dev/null +++ b/nemo_rl/data_plane/column_io.py @@ -0,0 +1,198 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Column-level helpers above :class:`DataPlaneClient`. + +These are thin wrappers around :meth:`get_samples` / :meth:`put_samples` +that operate on **columns** (named fields) of a partition ā€” not on the +driver process specifically. The driver uses them to fetch a slice and +materialize / write deltas back; worker-side dispatches use the +equivalents on ``AbstractPolicyWorker`` (``self._fetch(meta)`` / +``self._write_back``). + + * :func:`read_columns` ā€” ``get_samples + materialize`` (decode jagged + + object-array fields into a :class:`BatchedDataDict`). + * :func:`write_columns` ā€” pack-to-wire + ``put_samples`` for deltas + against an existing :class:`KVBatchMeta`. + * :func:`kv_first_write` ā€” pack-to-wire + ``put_samples`` for the + rollout-actor's first put of a partition. Returns a new + :class:`KVBatchMeta`. +""" + +from typing import Any, Sequence + +import numpy as np +import torch + +from nemo_rl.data.llm_message_utils import attach_message_log_view +from nemo_rl.data_plane.codec import materialize, pack_jagged_fields +from nemo_rl.data_plane.interfaces import DataPlaneClient, KVBatchMeta +from nemo_rl.data_plane.schema import GLOBAL_FORWARD_PAD_SEQLEN, Layout +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + +def round_up(value: int, multiple: int) -> int: + """Smallest ``multiple``-aligned int ā‰„ ``value`` (no-op when ``multiple <= 1``).""" + if multiple <= 1: + return value + return ((value + multiple - 1) // multiple) * multiple + + +def read_columns( + dp_client: DataPlaneClient, + meta: KVBatchMeta, + select_fields: Sequence[str], + *, + layout: Layout = "padded", + pad_value_dict: dict[str, Any] | None = None, +) -> BatchedDataDict[Any]: + """``get_samples(meta.sample_ids, select_fields=...) ā†’ materialize``. + + Pads to ``meta.extra_info[GLOBAL_FORWARD_PAD_SEQLEN]`` (minted on + the driver by ``TQPolicy._stamp_pad_seqlen`` and inherited by every + per-rank shard via :func:`shard_meta_for_dp`) ā€” so driver-fetched + and worker-returned columns land at one identical seq dim. + + Args: + dp_client: Data-plane client used for the underlying fetch. + meta: ``KVBatchMeta`` describing the keys to fetch. + select_fields: Fields to fetch. + layout: Materialization layout (``"padded"`` or ``"jagged"``). + pad_value_dict: Per-field pad value for jagged tensors (e.g. + ``input_ids ā†’ pad_token_id``); defaults to 0. + + Returns: + ``BatchedDataDict`` with the requested fields, materialized. + """ + td = dp_client.get_samples( + sample_ids=meta.sample_ids, + partition_id=meta.partition_id, + select_fields=list(select_fields), + ) + pad_to_seqlen = int((meta.extra_info or {}).get(GLOBAL_FORWARD_PAD_SEQLEN, 0)) + data = materialize( + td, + layout=layout, + pad_value_dict=pad_value_dict, + pad_to_seqlen=pad_to_seqlen, + ) + attach_message_log_view(data) + return data + + +def write_columns( + dp_client: DataPlaneClient, + meta: KVBatchMeta, + fields: "dict[str, torch.Tensor | np.ndarray]", +) -> None: + """``put_samples(meta.sample_ids, fields=...)``. + + Per-token tensor fields are converted to jagged via + :func:`pack_jagged_fields` so they land in TQ with the same row + lengths as the initial put. ``np.ndarray(dtype=object)`` leaves + pass through as-is. + + Args: + dp_client: Data-plane client used for the underlying put. + meta: ``KVBatchMeta`` describing the keys being written. + fields: Map of field name to tensor or object array. + """ + if not fields: + return + + seq_lens = meta.sequence_lengths + lengths = torch.tensor(seq_lens, dtype=torch.long) if seq_lens is not None else None + td = pack_jagged_fields(fields, lengths=lengths) + dp_client.put_samples( + sample_ids=meta.sample_ids, + partition_id=meta.partition_id, + fields=td, + ) + + +def kv_first_write( + final_batch_cpu: BatchedDataDict[Any], + *, + sample_ids: Sequence[str], + dp_client: DataPlaneClient, + partition_id: str, + extra_info: dict[str, Any] | None = None, + task_name: str = "train", + pad_to_multiple: int = 1, + tags: list[dict[str, Any]] | None = None, +) -> KVBatchMeta: + """Single flat ``put_samples`` of every tensor field in ``final_batch_cpu``. + + The rollout actor's first put of a partition. Caller mints + ``sample_ids`` (verl-style) ā€” the helper is rollout-shape-agnostic. + + Args: + final_batch_cpu: Rollout output already on CPU. Must contain + ``"sample_mask"`` (used as batch-size oracle: ``shape[0] == N``) + and ``"input_lengths"`` (per-row valid lengths for the jagged + pack). Tensor fields are packed jagged via + :func:`pack_jagged_fields`; ``np.ndarray(dtype=object)`` + leaves pass through. + sample_ids: Pre-minted per-sample ids, one per row of + ``final_batch_cpu``. + dp_client: Data-plane client used for the put. + partition_id: TQ partition to write into. + extra_info: Optional extra fields to attach to the returned meta. + task_name: Consumer task tag stamped on the returned meta. + pad_to_multiple: Seq-dim alignment recorded in ``extra_info`` so + readers pad to a multiple compatible with downstream backends + (mcore SP, PyTorch CP). + tags: Optional per-sample primitive metadata (one dict per row). + Stored on the TQ controller alongside the samples; travels + with ``KVBatchMeta`` through ``subset`` / ``concat`` / ``slice`` + so consumers can filter on it without fetching tensor data. + + Returns: + ``KVBatchMeta`` covering the written samples. + """ + n = int(final_batch_cpu["sample_mask"].shape[0]) + if n == 0 or len(sample_ids) != n: + raise ValueError( + f"kv_first_write: sample_ids ({len(sample_ids)}) must match batch size ({n})" + ) + if tags is not None and len(tags) != n: + raise ValueError( + f"kv_first_write: tags ({len(tags)}) must match batch size ({n})" + ) + lengths = final_batch_cpu["input_lengths"] + fields: dict[str, torch.Tensor | np.ndarray] = { + k: v + for k, v in final_batch_cpu.items() + if isinstance(v, torch.Tensor) + or (isinstance(v, np.ndarray) and v.dtype == object) + } + td = pack_jagged_fields(fields, lengths=lengths) + dp_client.put_samples( + sample_ids=list(sample_ids), + partition_id=partition_id, + fields=td, + tags=tags, + ) + + extras = dict(extra_info or {}) + if pad_to_multiple > 1: + extras["pad_to_multiple"] = int(pad_to_multiple) + return KVBatchMeta( + partition_id=partition_id, + task_name=task_name, + sample_ids=list(sample_ids), + fields=list(td.keys()), + sequence_lengths=[int(s) for s in lengths.tolist()], + extra_info=extras, + tags=[dict(t) for t in tags] if tags is not None else None, + ) diff --git a/nemo_rl/data_plane/docs/data-plane-async-proposal.md b/nemo_rl/data_plane/docs/data-plane-async-proposal.md new file mode 100644 index 00000000000..4b52bdd8e4b --- /dev/null +++ b/nemo_rl/data_plane/docs/data-plane-async-proposal.md @@ -0,0 +1,314 @@ + + +# Async path (proposed) + +The data-plane interface covers both sync and async, but the **sync +trainer (`grpo_train_sync`) uses only half of it**. The other half is +reserved for the async trainer (not yet landed). Everything below +documents the design proposal and open questions for that path. None +of it is wired into production today. + +## Sync vs Async at a glance + +| Concern | Sync (implemented) | Async (TODO) | +|---|---|---| +| **Who knows the keys?** | Driver ā€” `SyncRolloutActor` returns `KVBatchMeta` with `meta.keys` populated | TQ ā€” trainer doesn't know which samples are ready until it asks | +| **Data fetch API** | `kv_batch_get(meta.keys, ..., select_fields=[...])` ā€” direct by key | `claim_meta(...)` ā†’ `get_data(meta)` ā€” discover-then-fetch | +| **Consumer cursor?** | Not needed ā€” driver controls who reads what | `claim_meta` advances a per-task cursor; `check_consumption_status` confirms drain | +| **Step boundary** | `kv_clear(meta.keys)` at end of step | Same | + +In sync mode the driver always knows exactly which keys are in TQ +because it triggered every write. The task-mediated API +(`claim_meta` / `get_data` / `check_consumption_status`) is implemented +and tested but **not yet wired into any production codepath** ā€” it's +the future async-trainer's entry point. + +### Why two API surfaces? + +The deciding question is **"does the caller already know the keys?"** + +- **Yes** ā†’ use direct-by-key (`kv_batch_get`). The sync trainer is + always in this case: the rollout actor's return value carries + `meta.keys`. Cheapest path, no coordination. +- **No** ā†’ use task-mediated (`claim_meta` ā†’ `get_data`). The async + trainer is in this case: rollouts and training run concurrently, so + the trainer must ask TQ "what's ready for me to consume?" The + consumer cursor (`task_name`) prevents the same sample from being + claimed twice. + +verl follows the same split ā€” its `ReplayBuffer.sample()` returns a +`KVBatchMeta` from keys it tracks via `global_steps` tags, then fetches +via `kv_batch_get`. No `claim_meta` is used in verl's sync trainer +either. + +## Proposed E2E flow ā€” async GRPO + +In the async path, rollout and training run concurrently on separate +Ray actors. The trainer doesn't know which samples are ready ahead of +time, so it uses the task-mediated half of the API +(`claim_meta` / `get_data` / `check_consumption_status`) instead of +direct-by-key reads. + +``` +[PRODUCER ā€” continuous, never waits for trainer] +ā”Œā”€ AsyncTrajectoryCollector (Ray @remote) ā” +ā”‚ async_utils/trajectory_collector.py ā”‚ +ā”‚ Loop: ā”‚ +ā”‚ rollout ā†’ flatten ā†’ mask ā†’ prompt extract ā”‚ +ā”‚ kv_first_write(bulk, keys=[v_p_g, ā€¦]) ā”‚ +ā”‚ ā†’ dp_client.kv_batch_put ā”‚ +ā”‚ Pushes only KVBatchMeta onto an in-memory replay buffer ā”‚ +ā”‚ (bulk lives in TQ, never on the driver). ā”‚ +ā””ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”˜ + +[CONSUMER ā€” async trainer] +ā”Œā”€ DRIVER Ā· async grpo trainer (proposed) ā” +ā”‚ ā‘  policy.prepare_step(num_samples, group_size) ā”‚ +ā”‚ ā†’ register_partition("train", DP_TRAIN_FIELDS, ā”‚ +ā”‚ consumer_tasks=["prev_lp","ref_lp","train"])ā”‚ +ā”‚ ā”‚ +ā”‚ ā‘” meta = dp_client.claim_meta( ā”‚ +ā”‚ partition_id="train", ā”‚ +ā”‚ task_name="train", ā”‚ +ā”‚ required_fields=DP_TRAIN_FIELDS, ā”‚ +ā”‚ batch_size=GBS, ā”‚ +ā”‚ ) ā”‚ +ā”‚ ā†‘ BLOCKS until GBS samples have all required fields produced. ā”‚ +ā”‚ This is the *only* point where the per-task cursor advances ā€” ā”‚ +ā”‚ TQ's underlying ``get_meta(mode="fetch")`` marks those samples ā”‚ +ā”‚ as consumed by ``task_name``, so they won't be returned again ā”‚ +ā”‚ to the same task. ā”‚ +ā”‚ ā”‚ +ā”‚ ā‘¢ data = dp_client.get_data(meta, select_fields=ā€¦) ā”‚ +ā”‚ ā†‘ Pure key-list fetch (no cursor advancement here ā€” that already ā”‚ +ā”‚ happened at claim_meta). Or call ``policy.train_from_meta(meta)``ā”‚ +ā”‚ and let the workers fetch per-rank. ā”‚ +ā”‚ ā”‚ +ā”‚ ā‘£ training: same shard_meta_for_dp + fan-out as sync. ā”‚ +ā”‚ Workers fetch per-rank via dp_client.kv_batch_get and materialize. ā”‚ +ā”‚ ā”‚ +ā”‚ ā‘¤ Sync barrier before clearing: ā”‚ +ā”‚ dp_client.check_consumption_status( ā”‚ +ā”‚ "train", task_names=["prev_lp","ref_lp","train"]) ā”‚ +ā”‚ ā†‘ True iff every consumer task has drained ā€” safe to drop the data.ā”‚ +ā”‚ ā”‚ +ā”‚ ā‘„ dp_client.kv_clear(keys=meta.keys, partition_id="train") ā”‚ +ā””ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”˜ +``` + +**Why these methods are needed in async (but not sync):** + +| Method | Async role | Sync equivalent | +|---|---|---| +| `claim_meta` | discover + claim ready samples; per-task cursor prevents double-claim | not needed ā€” actor returns `meta.keys` directly | +| `get_data` | resolve meta ā†’ TensorDict (pure key-list fetch ā€” no cursor advancement) | not needed ā€” workers call `kv_batch_get` directly | +| `check_consumption_status` | safe-clear barrier when multiple consumers must drain before kv_clear | not needed ā€” single-thread Python ordering guarantees drain order | + +## Filtering without fetching bulk + +**Design constraint:** rollout writes samples continuously; many will +be discarded (off-policy beyond tolerance, DAPO `std == 0`, +format-check failures, length thresholds, ā€¦). The filter decision +**must not require reading bulk tensor data**. + +The filter state has to live somewhere small. Three alternative +options ā€” pick one based on what TQ/dataplane features are available +and how decoupled you want the cleanup to be. + +### Option 1 ā€” In TQ as a gating field (works today) + +The producer (or an intermediate stage) writes a small marker column +ONLY for samples that should be visible to downstream tasks. The +consumer `claim_meta(required_fields=["marker"])` only matches +samples where that field exists. + +```python +# Producer writes a small bool per survivor: +dp_client.kv_batch_put( + keys=survivor_keys, partition_id="train", + fields=TensorDict({"_train_ready": torch.ones(K)}, batch_size=[K]), +) +# Trainer never sees the non-survivors: +meta = dp_client.claim_meta(task_name="train", + required_fields=["input_ids", "_train_ready"], + batch_size=GBS) +``` + +- āœ… Server-side enforcement; consumer needs no special exclusion logic. +- āœ… Works with TQ as-is. +- āœ— Decision must be made at write time; no good story for filters + that become true *after* the write (e.g. weight-version drift). + +### Option 2 ā€” In TQ as tags (needs tag propagation in `KVBatchMeta`) + +The producer stamps primitive metadata (`weight_version`, `std`, +`total_reward`, `produced_at`) as **tags** on each key. Tags live on +the TQ controller alongside production status; reading them needs no +data RPC. The consumer inspects them in-memory: + +```python +# Producer: +tags = [{"weight_version": v, "std": s.item(), "produced_at": t} + for s, t in zip(stds, timestamps)] +dp_client.kv_batch_put(keys=keys, partition_id="train", fields=..., tags=tags) + +# Consumer (post-claim, no data fetch): +meta = dp_client.claim_meta(task_name="train", required_fields=[...], batch_size=K) +survivors = [i for i, tag in enumerate(meta.tags) + if current_version - tag["weight_version"] <= MAX_AGE] +meta = meta.subset(survivors) +``` + +- āœ… Zero data fetch ā€” tags travel with the meta. +- āœ… Works for *time-varying* filters (compare tag vs. current state). +- āœ— **Requires our `KVBatchMeta` to expose `tags`** (todo ā€” see + feature proposal below). + +### Option 3 ā€” Outside TQ entirely, in `AsyncTrajectoryCollector` + +The collector keeps a small driver-side ledger: `dict[key, +SampleMetadata]` tracking `weight_version`, `produced_at`, `status`, +etc. Sampling for training first consults the ledger, applies the +filter, and only then issues direct-by-key reads against TQ. TQ never +sees the filter ā€” it's just a KV store. + +```python +# inside AsyncTrajectoryCollector (Ray @remote) +def sample(self, batch_size: int, max_age: int) -> KVBatchMeta: + current_v = self._current_weight_version + survivor_keys = [ + k for k, m in self._ledger.items() + if (current_v - m.weight_version) <= max_age and m.status == "ready" + ][:batch_size] + return KVBatchMeta( + partition_id="train", task_name=None, + keys=survivor_keys, + fields=DP_TRAIN_FIELDS, + sequence_lengths=[self._ledger[k].seq_len for k in survivor_keys], + ) +``` + +- āœ… Zero TQ-side changes. +- āœ… Maximum flexibility ā€” any predicate, any state. +- āœ— Two sources of truth (collector ledger vs. TQ controller). On a + collector crash the ledger evaporates; needs reconciliation (e.g. + walk TQ partition on restart and reseed). + +## Timestamping / staleness specifically + +A common case worth singling out: rollouts produced under weight +version `v` may be too stale by version `v + N`. Four ways to handle +it, no bulk fetch needed in any of them: + +| Approach | Where state lives | Filter cost | Needs new feature? | +|---|---|---|---| +| Tag-stamp `weight_version`; consumer post-filters | TQ tags | zero | nemo-rl `KVBatchMeta.tags` propagation | +| Small `weight_version` field; `get_data(select_fields=["weight_version"])` | TQ field | one tiny RPC per claim | none | +| **Versioned partitions** (`train_v17`, `train_v18`, ā€¦) | TQ partition naming | zero | partition lifecycle helpers | +| `AsyncTrajectoryCollector` ledger with TTL | driver-side dict | zero | new collector method | + +**Versioned partitions** is interesting because it makes wholesale +staleness handling free: producers write into `train_v`, +trainer claims from `[train_v .. train_v]`, and +`kv_clear(partition_id="train_v")` retires an entire generation +of samples in one call. + +## Mark-as-stale, defer the kv_clear + +Filtered keys' bulk still sits in TQ. Two cleanup patterns: + +**Pattern A ā€” driver-side stale set + batched clear (recommended for +single-collector deployments):** + +```python +stale_keys: set[str] = set() +stale_keys.update(filter_meta.keys[i] for i in non_survivors) + +# Periodically (every K steps or size threshold): +if len(stale_keys) > 4096: + dp_client.kv_clear(keys=list(stale_keys), partition_id="train") + stale_keys.clear() +``` + +No TQ-side coordination. Bulk lingers briefly, bounded by the threshold. + +**Pattern B ā€” TQ-side stale-marker field + cleanup task (decoupled):** + +`claim_meta` filters on field production, not tag values ā€” so marking +via tags alone doesn't gate cleanup. Write a dedicated marker field: + +```python +dp_client.kv_batch_put( + keys=stale_keys, partition_id="train", + fields=TensorDict({"_stale": torch.ones(len(stale_keys), dtype=torch.bool)}, + batch_size=[len(stale_keys)]), +) +# A separate cleanup task: +cleanup_meta = dp_client.claim_meta( + partition_id="train", task_name="cleanup", + required_fields=["_stale"], batch_size=K, +) +dp_client.kv_clear(keys=cleanup_meta.keys, partition_id="train") +``` + +Pattern A is simpler. Pattern B decouples the cleanup cadence from +the filter site (useful if multiple producers can mark stale). + +## Proposed enhancements + +**TQ / data-plane side (in priority order):** + +1. **Propagate `tags` through nemo-rl `KVBatchMeta`** (small change, + high leverage). TQ's `KVBatchMeta` already carries `tags: + list[dict]`; our `interfaces.py:KVBatchMeta` only lifts + `input_lengths`. Add `tags: list[dict] | None` and have the + adapter pass them through. Unlocks Option 2 entirely. +2. **Server-side tag filtering in `claim_meta`**: e.g. + `claim_meta(..., tag_filter=lambda t: t["weight_version"] >= cutoff)`. + Today the consumer must claim everything ready and then filter + in-memory; a tag predicate would push this server-side. Requires + upstream TQ change. +3. **Versioned-partition helpers**: convenience methods + `register_versioned_partition(prefix, version)` + `claim_meta` + variant that takes a partition range. Cheap because TQ already + supports per-partition lifecycle. + +**`AsyncTrajectoryCollector` side (no TQ changes needed):** + +1. **Per-key ledger**: `dict[str, SampleMetadata]` on the collector + actor, populated at write time with `weight_version`, + `produced_at`, `seq_len`, `status`. +2. **`sample(batch_size, predicate)`**: returns a `KVBatchMeta` of + survivors after applying `predicate` to ledger entries. Trainer + never touches TQ for filtering. +3. **Mark-stale set + periodic batched `kv_clear`**: collector also + owns a background coroutine that drains stale keys on a cadence + (every K steps or by buffer pressure). +4. **Backpressure hook**: when ledger size approaches + `storage_capacity`, evict by oldest weight version. Decouples + producer from training rate. + +The collector-side path is the cheapest to land (zero TQ changes) and +gives the most flexibility; the TQ-side path scales better when +filtering needs to live close to the data (e.g. multiple trainers +filtering differently on the same partition). + +## Open questions + +- **`required_fields` granularity**: gate trainer on the full + `DP_TRAIN_FIELDS` set, or pipeline ā€” start training as soon as + `input_ids` + `generation_logprobs` are ready and gate on + `advantages` per microbatch? +- **Stale-data policy**: if the producer is multiple weight-versions + ahead of the trainer, drop those samples or use them with + importance-sampling correction? +- **Polling cadence**: `claim_meta_poll_interval_s` controls how often + `claim_meta` retries. Too aggressive = wasted CPU; too lazy = + trainer-rollout coupling. +- **Backpressure**: if rollout outpaces training, when does the + producer start blocking on TQ capacity? + (`storage_capacity` Ɨ `num_storage_units` is the hard cap.) +- **Cleanup cadence**: stale-key batch size for `kv_clear` ā€” + per-step, per-N-steps, or size-threshold? diff --git a/nemo_rl/data_plane/factory.py b/nemo_rl/data_plane/factory.py new file mode 100644 index 00000000000..86b5a944813 --- /dev/null +++ b/nemo_rl/data_plane/factory.py @@ -0,0 +1,67 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Single entrypoint that maps a :class:`DataPlaneConfig` to a client.""" + +from __future__ import annotations + +from nemo_rl.data_plane.interfaces import DataPlaneClient, DataPlaneConfig + + +def build_data_plane_client( + cfg: DataPlaneConfig | None, *, bootstrap: bool = True +) -> DataPlaneClient: + """Construct the configured data-plane client. + + Dispatches on ``cfg["impl"]``. Only ``"transfer_queue"`` ships today; + other adapters can be added behind this factory without touching + call sites. Raises if data_plane is disabled ā€” the legacy trainer + (``nemo_rl.algorithms.grpo.grpo_train``) should be used in that case + rather than a NoOp fallback here. + + Args: + cfg: Data-plane config; must have ``enabled=True``. + bootstrap: ``True`` on the driver ā€” bootstraps the TQ + controller. ``False`` on worker processes ā€” connects to the + existing controller (avoids creating a second named actor). + + Returns: + A configured ``DataPlaneClient``; wrapped in + :class:`MetricsDataPlaneClient` when observability is enabled. + """ + if cfg is None or not cfg["enabled"]: + raise ValueError( + "build_data_plane_client called with data_plane disabled. " + "Use the legacy nemo_rl.algorithms.grpo.grpo_train trainer " + "(which never engages the data plane) for that case." + ) + + impl = cfg["impl"] + if impl == "transfer_queue": + from nemo_rl.data_plane.adapters.transfer_queue import TQDataPlaneClient + + client: DataPlaneClient = TQDataPlaneClient(cfg, bootstrap=bootstrap) + else: + raise ValueError(f"unknown data_plane impl: {impl!r}") + + obs = cfg.get("observability") or {} + if obs.get("enabled", False): + from nemo_rl.data_plane.observability import ( + MetricsDataPlaneClient, + log_event, + ) + + on_event = obs.get("callback") or log_event + # pyrefly: obs.get returns Any, can't narrow to the expected callback type. + client = MetricsDataPlaneClient(client, on_event=on_event) # type: ignore[bad-argument-type] + return client diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py new file mode 100644 index 00000000000..6bdc5e940cf --- /dev/null +++ b/nemo_rl/data_plane/interfaces.py @@ -0,0 +1,422 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Stable boundary between NeMo-RL and data-plane implementations. + +Wire shape adapters must support: + * ``fields``: ``TensorDict`` with tensor leaves AND optional + ``NonTensorStack`` / ``NonTensorData`` leaves (TQ-native non-tensor + passthrough). TQ's storage backends handle encoding per backend + (simple keeps Python objects; mooncake_client pickles internally). + * ``tags``: ``list[dict[str, Any]]`` per-sample primitives (kept + separate from ``fields`` so non-tensor metadata like + ``input_lengths`` doesn't pollute the leaf-level schema). + * ``keys``: per-sample string uids. + * ``partition_id``: string-named address spaces with declared + ``consumer_tasks`` and ``fields`` schemas. + +All call sites in ``nemo_rl/algorithms``, ``nemo_rl/experience`` and +``nemo_rl/models`` go through :class:`DataPlaneClient` ā€” never +``import transfer_queue`` directly. This is what makes the +implementation swappable. + +See ``nemo_rl/data_plane/README.md`` for the full design. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any, Callable, Literal, NotRequired, Sequence, TypedDict + +from tensordict import TensorDict + + +class DataPlaneConfig(TypedDict): + """Feature-gated config; defaults to disabled. + + ``backend`` is the storage backend *inside* TransferQueue; it is owned by + the TQ adapter, not by NeMo-RL. ``impl`` selects which adapter we go + through. + + Required keys (always set in exemplar YAML ā€” never defaulted in code): + ``enabled``, ``impl``, ``backend``, ``storage_capacity``, + ``num_storage_units``, ``claim_meta_poll_interval_s``, + ``global_segment_size``, ``local_buffer_size``. + + ``global_segment_size`` / ``local_buffer_size`` are only *read* when + ``backend == "mooncake_cpu"``; the simple backend ignores them. + They are required (not NotRequired) so the YAML carries the full + schema and there are no hidden Python defaults. + """ + + enabled: bool + impl: Literal["transfer_queue"] + backend: Literal["simple", "mooncake_cpu"] + storage_capacity: int + num_storage_units: int + claim_meta_poll_interval_s: float + global_segment_size: int + local_buffer_size: int + controller_address: NotRequired[str] + ack_timeout_ms: NotRequired[int] + observability: NotRequired["ObservabilityConfig"] + + +class ObservabilityConfig(TypedDict): + """Optional middleware that records per-op metrics on the client. + + Off by default. When ``enabled=True`` the factory wraps the chosen + adapter with :class:`MetricsDataPlaneClient`. ``callback`` is + injected programmatically (callables don't round-trip through + YAML) ā€” set ``cfg["observability"]["callback"] = my_fn`` before + :func:`build_data_plane_client` to plug into wandb / file / log. + Default callback prints one line per op for debug. + """ + + enabled: bool + callback: NotRequired[Callable[[dict[str, Any]], None]] + + +@dataclass +class KVBatchMeta: + """Per-batch metadata for data-plane KV operations. + + Carries the per-sample IDs (``sample_ids``) that address rows in the + KV store plus per-row metadata (``fields``, ``sequence_lengths``, + ``tags``) needed for downstream routing without fetching tensor data. + Vocabulary is intentionally NeMo-RL-native rather than 1:1 with any + specific backend ā€” the adapter translates at the boundary. + + Two roles: + * Result type returned by :meth:`DataPlaneClient.claim_meta` ā€” callers + extract ``.sample_ids`` / ``.partition_id`` and pass them to + :meth:`get_samples` / :meth:`get_data`. + * Argument type for the per-DP-rank fetch entrypoints. + ``sequence_lengths`` lets the driver compute a balanced per-rank + shard from metadata only (control plane), without ever + materializing tensor data. + """ + + partition_id: str + task_name: str | None + sample_ids: list[str] + fields: list[str] | None = None + sequence_lengths: list[int] | None = None + extra_info: dict[str, Any] = field(default_factory=dict) + # Per-sample primitive sidecar. Aligned 1:1 with ``sample_ids`` when + # populated. Producers stamp filter scalars (std, total_reward, + # weight_version, ā€¦) here at ``put_samples`` time so consumers + # can filter without fetching tensor data. Mirrors verl's pattern + # and TQ's underlying ``KVBatchMeta.tags``. + tags: list[dict[str, Any]] | None = None + + def __post_init__(self) -> None: + if self.tags is not None and len(self.tags) != len(self.sample_ids): + raise ValueError( + f"KVBatchMeta: tags ({len(self.tags)}) must align 1:1 with " + f"sample_ids ({len(self.sample_ids)})" + ) + + @property + def size(self) -> int: + return len(self.sample_ids) + + def stamp_tags(self, scalars: dict[str, "Sequence[Any]"]) -> None: + """Mirror per-row scalar columns onto :attr:`tags`. + + Each entry in ``scalars`` is a length-``size`` sequence (list, + tensor, ndarray) whose elements are written to ``tags[i][name]``. + Initializes ``tags`` to a list of empty dicts if currently None. + """ + n = self.size + if self.tags is None: + self.tags = [{} for _ in range(n)] + for name, values in scalars.items(): + if len(values) != n: + raise ValueError( + f"stamp_tags: {name!r} has {len(values)} values, expected {n}" + ) + for i, v in enumerate(values): + self.tags[i][name] = v # type: ignore[bad-specialization] + + # ā”€ā”€ Pure-metadata transforms (no I/O) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + # Used by dynamic_sampling on the meta path: filter zero-std rows + # (subset), accumulate survivors across iterations (concat), trim + # an over-full cache to the training batch size (slice). Each + # returns a fresh KVBatchMeta ā€” caller is responsible for clear_samples- + # ing any uids dropped from the working set. + + def _replace( + self, + *, + sample_ids: list[str], + sequence_lengths: list[int] | None, + tags: list[dict[str, Any]] | None = None, + ) -> "KVBatchMeta": + """Return a copy with new sample_ids/sequence_lengths/tags, same metadata otherwise.""" + return KVBatchMeta( + partition_id=self.partition_id, + task_name=self.task_name, + sample_ids=list(sample_ids), + fields=self.fields, + sequence_lengths=list(sequence_lengths) + if sequence_lengths is not None + else None, + extra_info=dict(self.extra_info or {}), + tags=list(tags) if tags is not None else None, + ) + + def subset(self, indices: "Sequence[int]") -> "KVBatchMeta": + """Return a new meta with only the rows at ``indices`` (any order).""" + return self._replace( + sample_ids=[self.sample_ids[i] for i in indices], + sequence_lengths=( + [self.sequence_lengths[i] for i in indices] + if self.sequence_lengths is not None + else None + ), + tags=([self.tags[i] for i in indices] if self.tags is not None else None), + ) + + def slice(self, start: int, stop: int) -> "KVBatchMeta": + """Return a new meta with rows in the contiguous range ``[start, stop)``.""" + return self._replace( + sample_ids=self.sample_ids[start:stop], + sequence_lengths=( + self.sequence_lengths[start:stop] + if self.sequence_lengths is not None + else None + ), + tags=self.tags[start:stop] if self.tags is not None else None, + ) + + def concat(self, *others: "KVBatchMeta") -> "KVBatchMeta": + """Append ``others`` to ``self``. All metas must share ``partition_id``.""" + if any(o.partition_id != self.partition_id for o in others): + raise ValueError("KVBatchMeta.concat: partition_ids must match") + all_m = (self, *others) + sample_ids = [k for m in all_m for k in m.sample_ids] + all_have_lens = all(m.sequence_lengths is not None for m in all_m) + seq_lens = ( + [s for m in all_m for s in (m.sequence_lengths or [])] + if all_have_lens + else None + ) + all_have_tags = all(m.tags is not None for m in all_m) + tags = [t for m in all_m for t in (m.tags or [])] if all_have_tags else None + return self._replace( + sample_ids=sample_ids, sequence_lengths=seq_lens, tags=tags + ) + + +class DataPlaneClient(ABC): + """Stable, swappable data-plane boundary. + + The methods are split into three groups by intent. Argument order + mirrors the underlying ``transfer_queue`` API 1:1 so a future adapter + (e.g. ``nv-dataplane``) is a thin pass-through too. + + A. *Task-mediated* ā€” used by stages that wait for upstream production + via the per-task consumer counter: + :meth:`register_partition`, :meth:`claim_meta`, :meth:`get_data`, + :meth:`check_consumption_status`. + B. *Direct-by-key* ā€” used by stages that already know the exact uids + (e.g. driver-side fan-out to DP ranks): + :meth:`put_samples`, :meth:`get_samples`, :meth:`clear_samples`. + C. *Lifecycle* ā€” :meth:`close`. + + Stage-completion signal: there is intentionally no ``mark_consumed``. + The authoritative signal in TransferQueue is *field production* ā€” + when a stage calls :meth:`put_samples` for a new field, the controller + flips ``production_status[sample, field] = 1``. Downstream consumers + waiting on that field only see those samples once produced. + """ + + # ā”€ā”€ (A) task-mediated ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + @abstractmethod + def register_partition( + self, + partition_id: str, + fields: list[str], + num_samples: int, + consumer_tasks: list[str], + grpo_group_size: int | None = None, + enums: dict[str, list[str]] | None = None, + ) -> None: + """Declare the partition schema and consumer tasks. + + Args: + partition_id: Partition name. + fields: Superset of fields any producer may write here. + num_samples: Expected total samples; sizes controller arrays. + consumer_tasks: Named tasks; each gets its own consumption cursor. + grpo_group_size: Group size for GRPO balanced sampling. + enums: Per-field fixed-vocab string codec, shipped once at register. + """ + + @abstractmethod + def claim_meta( + self, + partition_id: str, + task_name: str, + required_fields: list[str], + batch_size: int, + dp_rank: int | None = None, + blocking: bool = True, + timeout_s: float = 60.0, + ) -> KVBatchMeta: + """Discover and **claim** up to ``batch_size`` ready samples. + + Advances ``task_name``'s per-sample consumption cursor (TQ's + ``mode='fetch'``); claimed uids won't be returned again. Samples + stay readable via :meth:`get_samples` until :meth:`clear_samples`. + + Args: + partition_id: Partition to claim from. + task_name: Consumer task whose cursor is advanced. + required_fields: Fields that must be produced for a sample to be claimable. + batch_size: Max samples to claim. + dp_rank: Reserved; driver-side balancing via :func:`shard_meta_for_dp` is used today. + blocking: Block until the batch can be claimed. + timeout_s: Max blocking time before raising. + + Returns: + ``KVBatchMeta`` for the claimed batch; pass to :meth:`get_data`. + """ + + @abstractmethod + def get_data( + self, + meta: KVBatchMeta, + select_fields: list[str] | None = None, + ) -> TensorDict: + """Resolve a meta to tensor data. + + Field-set resolution: (1) explicit ``select_fields``; (2) + ``meta.fields`` if non-None; (3) *fail loudly* ā€” never silently + fetch all fields. + + Args: + meta: From :meth:`claim_meta` or hand-built with explicit keys. + select_fields: Subset of fields to fetch. + + Returns: + ``TensorDict`` keyed by field name, batched along ``meta.sample_ids``. + """ + + @abstractmethod + def check_consumption_status( + self, partition_id: str, task_names: list[str] + ) -> bool: + """True iff every task has consumed all samples in the partition. + + Authoritative across workers ā€” uses TQ's controller-side counter, + not the per-process client cache. + + Args: + partition_id: Partition to check. + task_names: Tasks whose consumption cursors are inspected. + + Returns: + ``True`` iff every task in ``task_names`` has consumed all samples. + """ + + # ā”€ā”€ (B) direct-by-key (TQ-aligned signatures) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + @abstractmethod + def put_samples( + self, + sample_ids: list[str], + partition_id: str, + fields: TensorDict | None = None, + tags: list[dict[str, Any]] | None = None, + ) -> KVBatchMeta: + """Write fields for ``sample_ids`` ā€” the producer entrypoint. + + Writing a field flips the controller's ``production_status`` bit + for ``(sample, field)``; that flip is the "stage finished" signal + downstream consumers wait on. Tensor and ``NonTensorStack`` leaves + both pass through to TQ; non-tensor encoding is per-backend. + + Args: + sample_ids: Per-sample uids being written. + partition_id: Partition these samples belong to. + fields: Tensor / ``NonTensorStack`` leaves to write. + tags: Optional per-sample primitive metadata. + + Returns: + ``KVBatchMeta`` covering ``sample_ids`` ā€” usable for direct :meth:`get_samples`. + """ + + @abstractmethod + def get_samples( + self, + sample_ids: list[str], + partition_id: str, + select_fields: list[str], + ) -> TensorDict: + """Direct fetch by uids. + + Used by per-DP-rank slice fetches. Does NOT advance any per-task + consumption cursor ā€” that only happens via :meth:`claim_meta`. + + ``select_fields`` is required (no implicit "fetch every field" + fallback): bulk schemas are wide and silent over-fetch is the + most expensive shape the wire can take. Callers must name what + they read. + + Args: + sample_ids: Uids to fetch. + partition_id: Partition the samples live in. + select_fields: Subset of fields to fetch. + + Returns: + ``TensorDict`` keyed by field name, batched along ``sample_ids``. + """ + + @abstractmethod + def clear_samples( + self, + sample_ids: list[str] | None, + partition_id: str, + ) -> None: + """Drop key-value pairs. + + Explicit form (``sample_ids=[...]``) drops exactly those uids and + is the form callers should use whenever they have the meta in + hand ā€” both sync GRPO callers (driver passes ``meta.sample_ids``) + and future async-RL data-loader actors that don't share a + process-local registry with the producer. + + Convenience form (``sample_ids=None``) drops "everything this + process knows produced in this partition". Adapters implement + this via a local registry populated by :meth:`put_samples`, with + a fallback query to the underlying store. Useful for step-end + teardown when the caller is the producer (driver in sync GRPO). + Workers / loader actors that didn't produce the samples should + pass explicit IDs ā€” the ``None`` form may silently no-op for + them, and adapters are expected to warn when that happens. + + Args: + sample_ids: Uids to drop; ``None`` clears every uid this + process produced in the partition. + partition_id: Partition the samples live in. + """ + + # ā”€ā”€ (C) lifecycle ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + @abstractmethod + def close(self) -> None: + """Release controller / storage handles. Idempotent.""" diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py new file mode 100644 index 00000000000..63e551dc209 --- /dev/null +++ b/nemo_rl/data_plane/observability.py @@ -0,0 +1,345 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Lean per-op metrics decorator for ``DataPlaneClient``. + +Wraps any ``DataPlaneClient`` and invokes a single user-provided +callback on each operation. Each event is a flat dict:: + + {"op", "partition_id", "n_keys", "n_bytes", "wall_ms", "status"} + +Plug wandb / file logging / debug print at the call site by passing +``on_event=``. ``snapshot()`` returns cumulative +totals **plus** live memory consumption: ``bytes_outstanding`` (sum of +bytes currently held in TQ, i.e. put minus cleared) and +``peak_bytes_outstanding`` (high-water mark over the run lifetime). +""" + +from __future__ import annotations + +import logging +from dataclasses import asdict, dataclass +from time import monotonic +from typing import Any, Callable, Literal, TypedDict + +EventStatus = Literal["ok", "error", "timeout"] + + +class DataPlaneEvent(TypedDict): + op: str + partition_id: str + n_keys: int + n_bytes: int + wall_ms: float + status: EventStatus + + +import torch +from tensordict import TensorDict + +from nemo_rl.data_plane.interfaces import DataPlaneClient, KVBatchMeta + +logger = logging.getLogger(__name__) + + +def _td_bytes(td: TensorDict | None) -> int: + if td is None: + return 0 + total = 0 + for k in td.keys(include_nested=True, leaves_only=True): + v = td.get(k) + if not isinstance(v, torch.Tensor): + continue + t = v.values() if v.is_nested else v + total += t.numel() * t.element_size() + return total + + +def log_event(event: DataPlaneEvent) -> None: + logger.info("data_plane_event: %s", event) + + +@dataclass +class DataPlaneStats: + total_bytes: int = 0 + total_keys: int = 0 + total_ops: int = 0 + bytes_outstanding: int = 0 + peak_bytes_outstanding: int = 0 + # Anomaly trackers ā€” a wire-format regression that bloats bytes per + # row (cf. message_log view-aliasing pickle bug) shows up as a + # sudden spike in ``max_bytes_per_key_seen``. + max_bytes_per_key_seen: int = 0 + last_put_bytes_per_key: int = 0 + + +class MetricsDataPlaneClient(DataPlaneClient): + """Wrap a ``DataPlaneClient`` with a per-op callback hook.""" + + def __init__( + self, + inner: DataPlaneClient, + on_event: Callable[[DataPlaneEvent], None] | None = None, + ) -> None: + self._inner = inner + self._on_event = on_event or (lambda _: None) + self._stats = DataPlaneStats() + # Nested per-partition / per-key live byte counts. Populated on + # successful ``put_samples``; popped on successful ``clear_samples``. + # Bounded by the live key population, not cumulative traffic. + self._bytes_by_partition: dict[str, dict[str, int]] = {} + + def snapshot(self) -> dict[str, Any]: + """Return cumulative totals plus live byte / key outstanding counts.""" + out = asdict(self._stats) + out["n_keys_outstanding"] = sum( + len(d) for d in self._bytes_by_partition.values() + ) + return out + + def bytes_outstanding_by_partition(self) -> dict[str, int]: + """Per-partition breakdown of currently-held bytes.""" + return {p: sum(d.values()) for p, d in self._bytes_by_partition.items()} + + def _record_put(self, partition_id: str, keys: list[str], n_bytes: int) -> None: + """Attribute put bytes per key so a later ``clear_samples`` can subtract. + + Called after the underlying RPC succeeds so a failed put never + leaves the accounting inflated. + + Args: + partition_id: Partition the keys were written to. + keys: Per-sample uids that were written. + n_bytes: Total bytes written; distributed evenly across keys. + """ + if not keys or n_bytes <= 0: + return + per_key, remainder = divmod(n_bytes, len(keys)) + partition_dict = self._bytes_by_partition.setdefault(partition_id, {}) + for i, key in enumerate(keys): + share = per_key + (1 if i < remainder else 0) + partition_dict[key] = partition_dict.get(key, 0) + share + self._stats.bytes_outstanding += n_bytes + if self._stats.bytes_outstanding > self._stats.peak_bytes_outstanding: + self._stats.peak_bytes_outstanding = self._stats.bytes_outstanding + + def _record_clear(self, partition_id: str, keys: list[str] | None) -> None: + """Reverse the put accounting for ``keys``. + + Called after the underlying RPC succeeds so a failed clear keeps + the accounting consistent with TQ's actual state. + + Args: + partition_id: Partition the keys were dropped from. + keys: Uids dropped; ``None`` means the whole partition was cleared. + """ + partition_dict = self._bytes_by_partition.get(partition_id) + if partition_dict is None: + return + if keys is None: + freed = sum(partition_dict.values()) + del self._bytes_by_partition[partition_id] + else: + freed = 0 + for key in keys: + freed += partition_dict.pop(key, 0) + if not partition_dict: + del self._bytes_by_partition[partition_id] + self._stats.bytes_outstanding -= freed + + def _run( + self, + op: str, + partition_id: str, + fn: Callable[[], Any], + *, + n_keys: int = 0, + n_bytes: int = 0, + ) -> Any: + """Run ``fn`` and emit one observability event with wall-time and status. + + Args: + op: Operation tag (``"put"``, ``"get"``, ``"clear"``, etc.). + partition_id: Partition the op targets. + fn: Zero-arg callable that invokes the inner client. + n_keys: Key count if known up front; otherwise inferred from + the return value (``KVBatchMeta.sample_ids``). + n_bytes: Byte estimate; overridden by ``_td_bytes`` when the + return is a ``TensorDict``. + + Returns: + Whatever ``fn`` returned. + """ + t0 = monotonic() + try: + out = fn() + except TimeoutError: + self._emit(op, partition_id, n_keys, n_bytes, t0, "timeout") + raise + except Exception: + self._emit(op, partition_id, n_keys, n_bytes, t0, "error") + raise + # If the call returns a TensorDict, the read-side bytes are more + # informative than the input estimate. + if isinstance(out, TensorDict): + n_bytes = _td_bytes(out) + elif isinstance(out, KVBatchMeta) and not n_keys: + n_keys = len(out.sample_ids) + self._emit(op, partition_id, n_keys, n_bytes, t0, "ok") + return out + + def _emit( + self, + op: str, + partition_id: str, + n_keys: int, + n_bytes: int, + t0: float, + status: EventStatus, + ) -> None: + event: DataPlaneEvent = { + "op": op, + "partition_id": partition_id, + "n_keys": int(n_keys), + "n_bytes": int(n_bytes), + "wall_ms": (monotonic() - t0) * 1000.0, + "status": status, + } + self._on_event(event) + if status == "ok": + self._stats.total_bytes += n_bytes + self._stats.total_keys += n_keys + self._stats.total_ops += 1 + if op == "put" and n_keys: + per_key = n_bytes // n_keys + self._stats.last_put_bytes_per_key = per_key + if per_key > self._stats.max_bytes_per_key_seen: + self._stats.max_bytes_per_key_seen = per_key + + def register_partition( + self, + partition_id, + fields, + num_samples, + consumer_tasks, + grpo_group_size=None, + enums=None, + ): + self._run( + "register", + partition_id, + lambda: self._inner.register_partition( + partition_id, + fields, + num_samples, + consumer_tasks, + grpo_group_size=grpo_group_size, + enums=enums, + ), + n_keys=int(num_samples), + ) + + def claim_meta( + self, + partition_id, + task_name, + required_fields, + batch_size, + dp_rank=None, + blocking=True, + timeout_s=60.0, + ): + return self._run( + "claim_meta", + partition_id, + lambda: self._inner.claim_meta( + partition_id, + task_name, + required_fields, + batch_size, + dp_rank=dp_rank, + blocking=blocking, + timeout_s=timeout_s, + ), + ) + + def get_data(self, meta, select_fields=None): + return self._run( + "get_data", + meta.partition_id, + lambda: self._inner.get_data(meta, select_fields=select_fields), + n_keys=len(meta.sample_ids), + ) + + def check_consumption_status(self, partition_id, task_names): + return self._run( + "check_consumption_status", + partition_id, + lambda: self._inner.check_consumption_status(partition_id, task_names), + ) + + def put_samples(self, sample_ids, partition_id, fields=None, tags=None): + n_bytes = _td_bytes(fields) + # Materialize once: ``_run`` consumes its lambda and we also need + # to attribute bytes per sample after success. + sample_ids_list = ( + sample_ids if isinstance(sample_ids, list) else list(sample_ids) + ) + out = self._run( + "put", + partition_id, + lambda: self._inner.put_samples( + sample_ids_list, + partition_id, + fields=fields, + tags=tags, + ), + n_keys=len(sample_ids_list), + n_bytes=n_bytes, + ) + self._record_put(partition_id, sample_ids_list, n_bytes) + return out + + def get_samples(self, sample_ids, partition_id, select_fields): + return self._run( + "get", + partition_id, + lambda: self._inner.get_samples( + sample_ids, + partition_id, + select_fields=select_fields, + ), + n_keys=len(sample_ids), + ) + + def clear_samples(self, sample_ids, partition_id): + sample_ids_list = ( + sample_ids + if (sample_ids is None or isinstance(sample_ids, list)) + else list(sample_ids) + ) + n_keys = len(sample_ids_list) if sample_ids_list is not None else 0 + self._run( + "clear", + partition_id, + lambda: self._inner.clear_samples(sample_ids_list, partition_id), + n_keys=n_keys, + ) + self._record_clear(partition_id, sample_ids_list) + + def close(self) -> None: + self._run( + "close", + "", + lambda: self._inner.close(), + ) diff --git a/nemo_rl/data_plane/preshard.py b/nemo_rl/data_plane/preshard.py new file mode 100644 index 00000000000..f9ce2fdc6c7 --- /dev/null +++ b/nemo_rl/data_plane/preshard.py @@ -0,0 +1,182 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Driver-side balanced packing + per-rank fan-out helpers. + +Shared by sync and async data-plane trainers. Operates on full +``BatchedDataDict``s and relies on ``shard_by_batch_size``'s +``bin_count_multiple=DP_world`` behavior to keep per-rank microbatch +counts uniform ā€” without that, sequence packing / dynamic batching +produce variable per-rank bin counts and Megatron deadlocks at the +first cross-DP collective. +""" + +from __future__ import annotations + +from typing import Any, Optional + +import torch + +from nemo_rl.data_plane.interfaces import KVBatchMeta +from nemo_rl.data_plane.schema import ( + ELEM_COUNTS_PER_GB, + INPUT_IDS, + INPUT_LENGTHS, + META_IDX, + MICRO_BATCH_INDICES, + MICRO_BATCH_LENGTHS, + SAMPLE_MASK, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + +def shard_meta_for_dp( + meta: KVBatchMeta, + *, + dp_world: int, + batch_size: Optional[int] = None, + sequence_packing_args: Optional[dict[str, Any]] = None, + dynamic_batching_args: Optional[dict[str, Any]] = None, +) -> tuple[list[KVBatchMeta], Optional[list[int]]]: + """Pure key-list split: assign ``meta.sample_ids`` to ``dp_world`` ranks. + + Seq-len-aware on top of ``shard_by_batch_size``. No I/O, no key + minting. Used for every dispatch after rollout (logprob, ref-logprob, + train); the rollout actor's first write goes through + :func:`nemo_rl.experience.sync_rollout_actor.kv_first_write` directly. + + Per-rank packing metadata (``micro_batch_indices`` / + ``micro_batch_lengths`` / ``elem_counts_per_gb``) is set in each + shard's ``extra_info`` so the ``*_presharded`` worker can reattach + packing as it does on the legacy fan-out path. + + Args: + meta: Full-batch ``KVBatchMeta`` with ``sequence_lengths`` populated. + dp_world: Number of DP ranks. + batch_size: Total samples; ``None`` for the logprob path, GBS for train. + sequence_packing_args: Packing config dict for ``shard_by_batch_size``. + dynamic_batching_args: Dynamic-batching config dict; mutually exclusive with the above. + + Returns: + ``(per_rank_metas, unsorted_indices)``. ``unsorted_indices`` is + the inverse permutation that maps DP-rank-order outputs back to + original ``meta.sample_ids`` order (feed to + ``BatchedDataDict.reorder_data`` post-aggregation); ``None`` if + no reorder occurred. + """ + n = len(meta.sample_ids) + if n == 0: + raise ValueError("shard_meta_for_dp: empty meta ā€” nothing to shard") + if meta.sequence_lengths is None or len(meta.sequence_lengths) != n: + raise ValueError( + "shard_meta_for_dp requires meta.sequence_lengths populated and " + f"of length {n} (got {meta.sequence_lengths!r}). The rollout " + "actor's fan-out should populate this from input_lengths." + ) + if sequence_packing_args is not None and dynamic_batching_args is not None: + raise ValueError( + "Pass at most one of sequence_packing_args / dynamic_batching_args." + ) + + seq_lens = list(meta.sequence_lengths) + # Skeleton BatchedDataDict ā€” `shard_by_batch_size` only needs + # input_ids (placeholder), input_lengths (real), sample_mask (ones). + # ``meta_idx`` lets us recover which original meta index each shard row + # corresponds to, so we can slice ``meta.sample_ids`` per rank. + # + # ``INPUT_IDS`` seq dim sizing: the dynamic-batching microbatch planner + # in ``BatchedDataDict.shard_by_batch_size`` reads ``input_ids.shape[1]`` + # as an ``unpadded_seqlen`` cap (``min(padded_seqlen, unpadded_seqlen)``). + # A trivial ``(n, 1)`` shape made the cap clamp every microbatch length + # to 1, producing bogus ``micro_batch_lengths`` that, when consumed by + # workers, truncated real sequences to 1 token ā†’ zero grad_norm. Size + # the placeholder to ``max_tokens_per_microbatch`` (the largest seqlen + # the planner can ever request, per its own assertion) so the cap is + # never the binding factor. Memory cost is small (object only ā€” bytes + # never get filled with real data; just used for shape lookups). + input_ids_seqlen = 1 + if dynamic_batching_args is not None: + input_ids_seqlen = int(dynamic_batching_args["max_tokens_per_microbatch"]) + skeleton = BatchedDataDict( + { + INPUT_IDS: torch.zeros(n, input_ids_seqlen, dtype=torch.int64), + INPUT_LENGTHS: torch.tensor(seq_lens, dtype=torch.int64), + SAMPLE_MASK: torch.ones(n, dtype=torch.float32), + META_IDX: torch.arange(n, dtype=torch.int64), + } + ) + + if dynamic_batching_args is not None: + sharded, _ = skeleton.shard_by_batch_size( + dp_world, + batch_size=batch_size, + # pyrefly: ignore # bad-argument-type + dynamic_batching_args=dynamic_batching_args, + ) + elif sequence_packing_args is not None: + sharded, _ = skeleton.shard_by_batch_size( + dp_world, + batch_size=batch_size, + # pyrefly: ignore # bad-argument-type + sequence_packing_args=sequence_packing_args, + ) + else: + sharded = skeleton.shard_by_batch_size(dp_world, batch_size=batch_size) + + base_extra: dict[str, Any] = dict(meta.extra_info or {}) + out: list[KVBatchMeta] = [] + flat_idx: list[int] = [] + for shard in sharded: + # pyrefly: ignore # no-matching-overload + idx_list: list[int] = shard[META_IDX].tolist() + flat_idx.extend(idx_list) + rank_sample_ids = [meta.sample_ids[i] for i in idx_list] + rank_seqlens = [seq_lens[i] for i in idx_list] + rank_extra = dict(base_extra) + # Per-shard packing metadata ā€” set by ``shard_by_batch_size`` when + # sequence_packing or dynamic_batching is enabled. Workers' + # *_presharded paths look these up off ``meta.extra_info`` to avoid + # re-packing locally. Propagation is critical: local re-packing on + # different real per-rank data produces varying microbatch counts, + # which desynchronizes NCCL collectives across DP ranks and trips + # the Watchdog timeout. + for attr in ( + MICRO_BATCH_INDICES, + MICRO_BATCH_LENGTHS, + ELEM_COUNTS_PER_GB, + ): + val = getattr(shard, attr, None) + if val is not None: + rank_extra[attr] = val + out.append( + KVBatchMeta( + partition_id=meta.partition_id, + task_name=meta.task_name, + sample_ids=rank_sample_ids, + fields=meta.fields, + sequence_lengths=rank_seqlens, + extra_info=rank_extra, + ) + ) + + # Build inverse permutation: unsorted[orig_idx] = position_in_aggregated. + # When workers' results are concatenated in DP-rank order, row `j` of + # the aggregate corresponds to original index `flat_idx[j]`. To restore + # original meta.sample_ids order, the caller does aggregated.reorder_data( + # unsorted_indices) ā€” same contract as `_shard_for_logprob`. + unsorted: Optional[list[int]] = None + if flat_idx != list(range(n)): + unsorted = [0] * n + for new_pos, old_idx in enumerate(flat_idx): + unsorted[old_idx] = new_pos + return out, unsorted diff --git a/nemo_rl/data_plane/schema.py b/nemo_rl/data_plane/schema.py new file mode 100644 index 00000000000..5f14e3c0eeb --- /dev/null +++ b/nemo_rl/data_plane/schema.py @@ -0,0 +1,62 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Shared constants and type aliases for the data-plane meta contract.""" + +from typing import Literal + +# Materialization layout for `codec.materialize` / `read_columns` / worker fetch. +Layout = Literal["padded", "jagged"] + +# Per-shard packing metadata keys in `KVBatchMeta.extra_info`. +MICRO_BATCH_INDICES = "micro_batch_indices" +MICRO_BATCH_LENGTHS = "micro_batch_lengths" +ELEM_COUNTS_PER_GB = "elem_counts_per_gb" +GLOBAL_FORWARD_PAD_SEQLEN = "global_forward_pad_seqlen" + +# Skeleton field names from `shard_meta_for_dp`. +INPUT_IDS = "input_ids" +INPUT_LENGTHS = "input_lengths" +SAMPLE_MASK = "sample_mask" +META_IDX = "meta_idx" + +# Tensor fields in the train partition. Rollout writes the input +# subset on first put; later stages add prev_logprobs / +# reference_policy_logprobs (workers) and advantages (driver). +DP_TRAIN_FIELDS = ( + "input_ids", + "input_lengths", + "generation_logprobs", + "prev_logprobs", + "reference_policy_logprobs", + "advantages", + "token_mask", + "sample_mask", +) + +# Subset fetched by logprob / ref-logprob workers. +LP_SEED_FIELDS = ( + "input_ids", + "input_lengths", + "token_mask", + "sample_mask", +) + +# Fields requested for KV-scale calibration. Positive include-list: +# calibration only handles seq-dim tensor inputs, so we name them +# explicitly. Train-side deltas (logprobs/advantages/masks) and +# wire-only message-log bulk fields are skipped by virtue of not being +# in this list. ``multi_modal_inputs`` covers VLM extras (pixel values, +# grid metadata, etc.) when present; it's harmlessly absent for +# text-only models so the filter skips it on those. +DP_CALIB_INPUT_FIELDS = (INPUT_IDS, INPUT_LENGTHS, "multi_modal_inputs") diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py new file mode 100644 index 00000000000..bd558e0f06b --- /dev/null +++ b/nemo_rl/data_plane/worker_mixin.py @@ -0,0 +1,507 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""TransferQueue awareness for policy workers, isolated from the base class. + +Mix into a worker class to add per-rank TQ-mediated entrypoints +(:meth:`train_presharded`, :meth:`get_logprobs_presharded`, +:meth:`get_reference_policy_logprobs_presharded`) without touching +``BasePolicyWorker``. Subclasses that don't need TQ keep their bare +inheritance and stay zero-cost. + +Subclasses must implement :meth:`_get_replica_group` (returns the +NCCL group of TPƗCPƗPP siblings within this DP rank, or ``None`` for +TP=CP=PP=1) and inherit ``train`` / ``get_logprobs`` / +``get_reference_policy_logprobs`` from the worker base. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal, Optional + +import torch + +FetchPolicy = Literal["auto", "independent", "leader_broadcast"] + +from nemo_rl.data.llm_message_utils import attach_message_log_view +from nemo_rl.data_plane.schema import ( + ELEM_COUNTS_PER_GB, + GLOBAL_FORWARD_PAD_SEQLEN, + MICRO_BATCH_INDICES, + MICRO_BATCH_LENGTHS, + Layout, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.utils.nsys import wrap_with_nvtx_name + +if TYPE_CHECKING: + from nemo_rl.data_plane import DataPlaneConfig, KVBatchMeta + from nemo_rl.data_plane.interfaces import DataPlaneClient + + +def _broadcast_batched_data_dict( + data: Optional[BatchedDataDict[Any]], + *, + is_leader: bool, + src: int, + group: Any, +) -> BatchedDataDict[Any]: + """Broadcast a BatchedDataDict from ``src`` to all ranks in ``group``. + + Two-phase to avoid pickling tensor payloads on the hot path: a small + descriptor (per-key dtype/shape) ships via ``broadcast_object_list`` + first, then each tensor's data ships via ``broadcast`` on its + current device. The leader supplies ``data``; non-leaders pass + ``None`` and get an empty BatchedDataDict filled in-place. + """ + # NCCL groups can only broadcast CUDA tensors; pick the broadcast + # device from the group backend so CPU TQ outputs are moved to GPU + # before NCCL broadcast. + backend = torch.distributed.get_backend(group) + bcast_device: Any = torch.cuda.current_device() if backend == "nccl" else "cpu" + + if is_leader: + assert data is not None, "leader must provide non-None data" + descriptor: list[Any] = [] + for k, v in data.items(): + if isinstance(v, torch.Tensor): + descriptor.append( + (k, "tensor", str(v.dtype), tuple(v.shape), str(v.device)) + ) + else: + descriptor.append((k, "raw", v)) + payload: list[Any] = [descriptor] + else: + payload = [None] + + torch.distributed.broadcast_object_list(payload, src=src, group=group) + descriptor = payload[0] + assert descriptor is not None + + # pyrefly: ignore # bad-assignment + out: BatchedDataDict[Any] = data if is_leader else BatchedDataDict() + for entry in descriptor: + key = entry[0] + kind = entry[1] + if kind == "tensor": + dtype_str, shape, src_device = entry[2], entry[3], entry[4] + if is_leader: + tensor = out[key] + if tensor.device.type != torch.device(bcast_device).type: + tensor = tensor.to(bcast_device) + out[key] = tensor + else: + dtype = getattr(torch, dtype_str.split(".")[-1]) + tensor = torch.empty(shape, dtype=dtype, device=bcast_device) + out[key] = tensor + torch.distributed.broadcast(tensor, src=src, group=group) + # Restore non-leader tensors to the leader's source device + # so downstream code sees the same layout pre-broadcast. + if ( + not is_leader + and torch.device(src_device).type != torch.device(bcast_device).type + ): + out[key] = tensor.to(src_device) + else: + if not is_leader: + out[key] = entry[2] + return out + + +class TQWorkerMixin: + """Adds TransferQueue per-rank fetch/write-back to a policy worker. + + The driver-side ``TQPolicy`` fans out per-rank ``KVBatchMeta``; + each worker calls ``self._fetch(meta, ...)`` to pull its slice from + TQ and runs the existing per-rank method body. + """ + + _dp_client: Optional[DataPlaneClient] = None + + def setup_data_plane(self, cfg: DataPlaneConfig) -> None: + """Connect this worker process's client to the existing TQ controller. + + Called once by the driver after worker construction. Idempotent. + """ + if self._dp_client is not None: + return + from nemo_rl.data_plane import build_data_plane_client + + # bootstrap=False ā€” the driver already created the named + # controller actor; this process attaches as a client. + self._dp_client = build_data_plane_client(cfg, bootstrap=False) + + def _require_dp_client(self) -> DataPlaneClient: + if self._dp_client is None: + raise RuntimeError( + "Data-plane client not initialised on worker. The driver " + "must call setup_data_plane(cfg) before invoking any " + "*_presharded entrypoint." + ) + return self._dp_client + + def _get_replica_group(self) -> Optional[Any]: + """NCCL group of TPƗCPƗPP siblings within this DP rank. + + ``None`` means "no siblings" (TP=CP=PP=1). Subclasses must + override using their parallelism state (DTensor ``device_mesh``, + Megatron ``parallel_state``). Returning ``None`` makes + :meth:`_fetch` use independent fetch; returning a group makes + it use leader-fetch + NCCL broadcast. + """ + return None + + def _pad_value_dict(self) -> dict[str, Any]: + """Per-field pad value used by :func:`materialize` to detile the jagged wire format. + + Token-id fields use the tokenizer pad id. + """ + pad_id = getattr(getattr(self, "tokenizer", None), "pad_token_id", None) + if pad_id is None: + return {} + return {"input_ids": pad_id, "prompt_ids_for_adv": pad_id} + + def _forward_pad_seqlen(self, meta: "KVBatchMeta") -> int: + """Cross-DP forward pad target, minted by :meth:`TQPolicy._stamp_pad_seqlen`.""" + return int((meta.extra_info or {}).get(GLOBAL_FORWARD_PAD_SEQLEN, 0)) + + def _fetch( + self, + meta: "KVBatchMeta", + *, + layout: Layout = "padded", + fetch_policy: FetchPolicy = "auto", + preprocess: Optional[Any] = None, + dp_aligned_seq_len: bool = True, + ) -> BatchedDataDict[Any]: + """Fetch this rank's slice from TQ and return a BatchedDataDict. + + Args: + meta: Per-rank ``KVBatchMeta`` from :func:`shard_meta_for_dp`. + Forward-pass pad target is read from + ``meta.extra_info[GLOBAL_FORWARD_PAD_SEQLEN]`` minted by + :meth:`TQPolicy._stamp_pad_seqlen`. + layout: Materialization layout (``"padded"`` or ``"jagged"``). + fetch_policy: ``"auto"`` uses leader-fetch + NCCL broadcast when + :meth:`_get_replica_group` returns a group, else independent + fetch (cheapest for TP=CP=PP=1). ``"independent"`` forces + every sibling to fetch. ``"leader_broadcast"`` forces the + broadcast path and asserts a replica group exists. + preprocess: Optional ``(worker, td) -> td`` applied between + materialize and return. + dp_aligned_seq_len: When True (default), right-pad the seq + dim for the forward pass. Disabled in tests that want + to observe per-rank local-pad behavior. + + Returns: + ``BatchedDataDict`` of this rank's slice. + """ + if fetch_policy not in {"auto", "independent", "leader_broadcast"}: + raise ValueError(f"unknown fetch_policy: {fetch_policy!r}") + + from nemo_rl.data_plane import materialize + + pad_value_dict = self._pad_value_dict() + replica_group = ( + self._get_replica_group() + if fetch_policy in {"auto", "leader_broadcast"} + else None + ) + if fetch_policy == "leader_broadcast" and replica_group is None: + raise RuntimeError( + "_fetch(fetch_policy='leader_broadcast') requires a " + "replica group, but _get_replica_group() returned None." + ) + + pad_to_seqlen = self._forward_pad_seqlen(meta) if dp_aligned_seq_len else 0 + + if replica_group is not None and replica_group.size() > 1: + is_leader = self._is_replica_leader() + leader = torch.distributed.get_global_rank(replica_group, 0) + if is_leader: + td = self._require_dp_client().get_samples( + sample_ids=meta.sample_ids, + partition_id=meta.partition_id, + select_fields=list(meta.fields), # type: ignore[no-matching-overload] + ) + data = materialize( + td, + layout=layout, + pad_value_dict=pad_value_dict, + pad_to_seqlen=pad_to_seqlen, + ) + else: + data = None + data = _broadcast_batched_data_dict( + data, + is_leader=is_leader, + src=leader, + group=replica_group, + ) + # Reconstruct message_log after broadcast so the views alias + # the per-rank local ``input_ids`` rather than the leader's. + attach_message_log_view(data) + if preprocess is not None: + data = preprocess(self, data) + return data + + td = self._require_dp_client().get_samples( + sample_ids=meta.sample_ids, + partition_id=meta.partition_id, + select_fields=list(meta.fields), # type: ignore[no-matching-overload] + ) + data = materialize( + td, + layout=layout, + pad_value_dict=pad_value_dict, + pad_to_seqlen=pad_to_seqlen, + ) + attach_message_log_view(data) + if preprocess is not None: + data = preprocess(self, data) + return data + + def _apply_packing_prep(self, data: BatchedDataDict[Any]) -> BatchedDataDict[Any]: + """Re-derive ``micro_batch_indices`` / ``micro_batch_lengths`` on the local slice. + + Uses ``shard_by_batch_size(shards=1, ...)``. The legacy DP path computes those + as a side effect of the DP-shard call; the TQ presharded path receives a + per-rank slice without them set, so we recompute here using ``self.cfg``. + """ + cfg = getattr(self, "cfg", None) + if not isinstance(cfg, dict): + return data + seqpack = cfg.get("sequence_packing", {}) or {} + dynbatch = cfg.get("dynamic_batching", {}) or {} + + if seqpack.get("enabled", False): + spa = { + "algorithm": seqpack["algorithm"], + "input_key": "input_ids", + "input_lengths_key": "input_lengths", + "sequence_length_pad_multiple": cfg[ + "make_sequence_length_divisible_by" + ], + "max_tokens_per_microbatch": seqpack["train_mb_tokens"], + } + packed, _ = data.shard_by_batch_size( + shards=1, + batch_size=None, + # pyrefly: ignore # bad-argument-type + sequence_packing_args=spa, + ) + return packed[0] + + if dynbatch.get("enabled", False): + dba = { + "input_key": "input_ids", + "input_lengths_key": "input_lengths", + "sequence_length_round": dynbatch["sequence_length_round"], + "max_tokens_per_microbatch": dynbatch["train_mb_tokens"], + } + sharded, _ = data.shard_by_batch_size( + shards=1, + batch_size=None, + # pyrefly: ignore # bad-argument-type + dynamic_batching_args=dba, + ) + return sharded[0] + + return data + + def _attach_or_repack_pack_metadata( + self, + data: BatchedDataDict[Any], + meta: "KVBatchMeta", + ) -> BatchedDataDict[Any]: + """Trust driver-supplied packing metadata or re-derive locally. + + When the driver pre-balanced packing across DP ranks it ships + ``micro_batch_indices`` / ``micro_batch_lengths`` (and optionally + ``elem_counts_per_gb``) in ``meta.extra_info``. Locally + re-packing produces variable bin counts across DP groups and + desyncs Megatron's per-microbatch collectives ā€” trust the driver + when it provided the metadata. + """ + extra = meta.extra_info or {} + if MICRO_BATCH_INDICES in extra and MICRO_BATCH_LENGTHS in extra: + data.micro_batch_indices = extra[MICRO_BATCH_INDICES] + data.micro_batch_lengths = extra[MICRO_BATCH_LENGTHS] + if ELEM_COUNTS_PER_GB in extra: + data.elem_counts_per_gb = extra[ELEM_COUNTS_PER_GB] + return data + return self._apply_packing_prep(data) + + def _local_coords(self) -> dict[str, int]: + """This worker's (axis -> local-rank) mapping. + + Subclasses MUST override: DTensor reads ``device_mesh``, + Megatron reads ``parallel_state``. There's no honest default ā€” + a missing impl would silently make every rank a writeback + leader and re-create the ``-601 ILLEGAL_CLIENT`` duplicate-write + bug. + """ + raise NotImplementedError( + f"{type(self).__name__} must implement _local_coords() to gate TQ writeback. " + "Return (axis -> local rank) from the worker's parallelism state." + ) + + def _is_replica_leader(self) -> bool: + """True iff this rank should perform per-DP-rank-unique side-effects. + + Examples include TQ write-back. Shares the same predicate the + driver uses to gate dispatch (:meth:`NamedSharding.is_axis_zero`) + ā€” fed by per-worker :meth:`_local_coords` instead of + ``NamedSharding.get_worker_coords``; same answer either way. + """ + from nemo_rl.distributed.named_sharding import REPLICATED_AXES, NamedSharding + + return NamedSharding.is_axis_zero(self._local_coords(), REPLICATED_AXES) + + def _write_back( + self, + meta: "KVBatchMeta", + fields: dict[str, torch.Tensor], + ) -> None: + """Leader-only ``put_samples(meta.sample_ids, fields=...)``. + + Per-token fields are jagged-packed via :func:`maybe_pack_jagged` + so they land with the same row lengths as the initial put; + without this a worker write-back (rectangular ``[N, S]``) would + mismatch the jagged ``input_ids`` on the next read. + + Args: + meta: Per-rank ``KVBatchMeta`` for this slice. + fields: Map of field name to tensor to write back. + """ + if not self._is_replica_leader() or not fields: + return + from nemo_rl.data_plane.column_io import write_columns + + write_columns(self._require_dp_client(), meta, fields) + + def _write_back_result_field( + self, + meta: "KVBatchMeta", + result: Any, + *, + result_key: str, + tq_field: str, + ) -> None: + """Single chokepoint for ``*_presharded`` write-backs. + + ``result`` is checked via the ``Mapping`` ABC because + ``BatchedDataDict`` is a ``UserDict`` (not ``dict``). + + Args: + meta: Per-rank ``KVBatchMeta`` for this slice. + result: Worker output containing ``result_key``. + result_key: Key into ``result`` for the tensor to write back. + tq_field: Field name on the TQ side. + """ + if self._dp_client is None: + return + from collections.abc import Mapping + + if not isinstance(result, Mapping) or result_key not in result: + raise RuntimeError( + f"_write_back_result_field: result type {type(result).__name__} " + f"missing key {result_key!r}; cannot write back." + ) + val = result[result_key] + if not isinstance(val, torch.Tensor): + raise TypeError( + f"_write_back_result_field: result[{result_key!r}] is " + f"{type(val).__name__}, expected torch.Tensor." + ) + if val.shape[0] != len(meta.sample_ids): + raise ValueError( + f"_write_back_result_field: shape mismatch ā€” " + f"result[{result_key!r}] has batch dim {val.shape[0]} " + f"but meta.sample_ids has {len(meta.sample_ids)}." + ) + self._write_back(meta, {tq_field: val.detach().to("cpu")}) + + @wrap_with_nvtx_name("policy_worker/train_presharded") + def train_presharded( + self, + meta: "KVBatchMeta", + loss_fn: Any, + eval_mode: bool = False, + gbs: Optional[int] = None, + mbs: Optional[int] = None, + ) -> dict[str, Any]: + """Per-rank training entrypoint. Fetch ā†’ packing prep ā†’ delegate.""" + data = self._fetch(meta) + data = self._attach_or_repack_pack_metadata(data, meta) + return self.train( # type: ignore[attr-defined] + data, + loss_fn=loss_fn, + eval_mode=eval_mode, + gbs=gbs, + mbs=mbs, + ) + + @wrap_with_nvtx_name("policy_worker/get_logprobs_presharded") + def get_logprobs_presharded( + self, + meta: "KVBatchMeta", + micro_batch_size: Optional[int] = None, + ) -> None: + """Per-rank logprob entrypoint. Fetch ā†’ packing prep ā†’ run ā†’ write back. + + Returns ``None`` ā€” the per-token tensor is committed to TQ via + :meth:`_write_back_result_field` under ``prev_logprobs``. + Callers fetch it through :meth:`TQPolicy.read_from_dataplane` ā€” + skipping the Ray plasma roundtrip on the (B, S) tensor. + ``del result`` drops the local reference before returning so the + worker doesn't carry the tensor into the next dispatch. + """ + data = self._fetch(meta) + data = self._attach_or_repack_pack_metadata(data, meta) + result: BatchedDataDict[Any] = self.get_logprobs( # type: ignore[attr-defined] + data=data, + micro_batch_size=micro_batch_size, + ) + self._write_back_result_field( + meta, + result, + result_key="logprobs", + tq_field="prev_logprobs", + ) + del result + + @wrap_with_nvtx_name("policy_worker/get_reference_policy_logprobs_presharded") + def get_reference_policy_logprobs_presharded( + self, + meta: "KVBatchMeta", + micro_batch_size: Optional[int] = None, + ) -> None: + """Per-rank reference-policy logprob entrypoint. + + See :meth:`get_logprobs_presharded` for the contract. Tensor + lives in TQ under ``reference_policy_logprobs``. + """ + data = self._fetch(meta) + data = self._attach_or_repack_pack_metadata(data, meta) + result: BatchedDataDict[Any] = self.get_reference_policy_logprobs( # type: ignore[attr-defined] + data=data, + micro_batch_size=micro_batch_size, + ) + self._write_back_result_field( + meta, + result, + result_key="reference_logprobs", + tq_field="reference_policy_logprobs", + ) + del result diff --git a/nemo_rl/distributed/named_sharding.py b/nemo_rl/distributed/named_sharding.py index 8225c9380a6..234a8094e30 100644 --- a/nemo_rl/distributed/named_sharding.py +++ b/nemo_rl/distributed/named_sharding.py @@ -15,6 +15,18 @@ import numpy as np +# Canonical axis names that get *replicated* (every rank holds the same +# data along these axes). Used as the default ``axes`` arg to +# :meth:`NamedSharding.is_axis_zero` for leader-rank gating ā€” the +# leader is the worker at coord 0 on every replicated axis. Keep this +# list as the single source of truth; a typo in a caller's inline list +# would silently route around the leader gate. +REPLICATED_AXES: tuple[str, ...] = ( + "tensor_parallel", + "context_parallel", + "pipeline_parallel", +) + class NamedSharding: """Represents an N-dimensional arrangement of ranks with named axes, facilitating data sharding, replication, and collection based on these axes. @@ -121,6 +133,17 @@ def get_worker_coords(self, worker_id: int) -> dict[str, int]: coords[axis_name] = indices[i].item() return coords + @staticmethod + def is_axis_zero(coords: dict[str, int], axes: Sequence[str]) -> bool: + """Returns True when ``coords`` has value 0 on every ``axes`` entry. + + Shared leader-rank check fed by ``TQWorkerMixin._local_coords`` + on the worker side; driver-side callers can pair with + ``get_worker_coords`` directly. Axes missing from ``coords`` are + treated as rank 0. + """ + return all(coords.get(ax, 0) == 0 for ax in axes) + def get_ranks_by_coord(self, **coords: int) -> list[int]: """Gets all ranks that match the specified coordinates for named axes. diff --git a/nemo_rl/distributed/ray_actor_environment_registry.py b/nemo_rl/distributed/ray_actor_environment_registry.py index 30b0ae80bd1..0b1f0edfe4f 100644 --- a/nemo_rl/distributed/ray_actor_environment_registry.py +++ b/nemo_rl/distributed/ray_actor_environment_registry.py @@ -45,6 +45,13 @@ "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector": PY_EXECUTABLES.VLLM, # ReplayBuffer needs vLLM environment to handle trajectory data from VllmGenerationWorker "nemo_rl.algorithms.async_utils.ReplayBuffer": PY_EXECUTABLES.VLLM, + # SyncRolloutActor doesn't import vllm directly ā€” policy_generation is a + # Ray actor handle. The VLLM env is needed because (1) transfer_queue is + # bundled into the VLLM venv (and the policy training venvs), and the + # actor writes flattened tensors to TQ via dp_client.put_samples; + # (2) same-node colocation with VllmGenerationWorker avoids duplicate + # venv caches. + "nemo_rl.experience.sync_rollout_actor.SyncRolloutActor": PY_EXECUTABLES.VLLM, "nemo_rl.environments.tools.retriever.RAGEnvironment": PY_EXECUTABLES.SYSTEM, "nemo_rl.environments.nemo_gym.NemoGym": PY_EXECUTABLES.NEMO_GYM, } diff --git a/nemo_rl/experience/sync_rollout_actor.py b/nemo_rl/experience/sync_rollout_actor.py new file mode 100644 index 00000000000..0b0034d5152 --- /dev/null +++ b/nemo_rl/experience/sync_rollout_actor.py @@ -0,0 +1,372 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Sync GRPO rollout actor ā€” sibling of ``async_utils``. + +Houses :class:`SyncRolloutActor`, the Ray actor that owns the multi-turn +rollout loop AND the post-rollout flatten / mask / prompt extraction / +reward shaping / baseline-std for a sync GRPO step. The driver dispatches +a per-step prompt batch + uids; the actor runs ``run_multi_turn_rollout`` +(or async / nemo_gym variants), then writes the bulk schema to TQ via +:func:`nemo_rl.data_plane.column_io.kv_first_write`. Only a ``KVBatchMeta`` +and a small per-sample ``driver_carry`` dict (rewards, masks, lengths, +baseline/std, prompt_ids_for_adv) cross back to the driver via Ray. + +**Goal ā€” rollout 1-hop put**: bulk tensors (input_ids, output_ids, +attention_mask, position_ids, multi_modal_inputs, generation_logprobs, +token_mask) stay actor-side until ``put_samples``, then live only in +TQ. Driver never holds these bytes between rollout finish and train +fan-out. + +The actor is the sync counterpart to +:class:`nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector`. It +intentionally does not buffer or stream ā€” sync GRPO consumes the whole +step batch in one call. +""" + +from __future__ import annotations + +import uuid +from typing import Any, Optional + +import numpy as np +import ray +import torch + +from nemo_rl.data_plane.column_io import kv_first_write +from nemo_rl.data_plane.interfaces import KVBatchMeta +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.environments.interfaces import EnvironmentInterface +from nemo_rl.experience.rollouts import ( + run_async_multi_turn_rollout, + run_async_nemo_gym_rollout, + run_multi_turn_rollout, +) +from nemo_rl.models.generation.interfaces import GenerationInterface + +# Carry keys producible by the rollout actor only when the caller opts in. +# These are np.ndarray(object) per-row arrays from decompose_message_log; the +# default driver_carry omits them because BatchedDataDict.select_indices on +# the training/dynamic-sampling path only handles tensors/lists. Validation +# requests them explicitly to print per-sample message logs. +OPT_IN_CARRY_KEYS: tuple[str, ...] = ("turn_roles", "turn_contents") + + +@ray.remote # pragma: no cover +class SyncRolloutActor: + """Per-step rollout dispatcher. + + Runs: rollout + flatten + mask + prompt extraction + baseline/std + TQ put. + Returns ``(meta, driver_carry, rollout_metrics, gen_metrics)``. + + Lifecycle: one instance per ``grpo_train_sync`` invocation. The driver + instantiates with the same handles it would normally pass to + ``run_multi_turn_rollout`` plus the data-plane config so the actor + can attach as a TQ client (``bootstrap=False`` ā€” controller is + bootstrapped on the driver via ``TQPolicy``). + """ + + def __init__( + self, + policy_generation: GenerationInterface, + tokenizer: Any, + task_to_env: dict[str, EnvironmentInterface], + master_config: Any, + dp_cfg: dict[str, Any], + ) -> None: + self.policy_generation = policy_generation + self.tokenizer = tokenizer + self.task_to_env = task_to_env + self.master_config = master_config + + from nemo_rl.data_plane import build_data_plane_client + + self._dp_client = build_data_plane_client(dp_cfg, bootstrap=False) + + def rollout_to_tq( + self, + input_batch: BatchedDataDict[Any], + *, + partition_id: str, + group_size: int = 1, + first_iter: bool = True, + finish_generation: bool = True, + task_to_env_override: Optional[dict[str, EnvironmentInterface]] = None, + carry_keys: Optional[list[str]] = None, + ) -> tuple[ + KVBatchMeta, + dict[str, Any], + dict[str, Any], + Optional[dict[str, Any]], + ]: + """Run the full per-step generation cycle and write bulk data to TQ. + + Bundles six steps into one Ray round-trip so the driver only sees + a single RPC instead of separate calls for each: + + 1. **Reset metrics** ā€” ``policy_generation.clear_logger_metrics()`` + clears per-step generation accumulators before the rollout. + 2. **Rollout** ā€” runs ``run_multi_turn_rollout`` (or the async / + nemo-gym variants) to produce ``final_batch``. + 3. **Flatten + mask + prompt extraction** ā€” converts + ``message_log`` layout to flat tensors; builds token mask, + sample mask, prompt-only ids, baseline/std. + 4. **Write bulk to TQ** ā€” ``kv_first_write`` puts every tensor + field in one flat ``put_samples``; the driver never touches + bulk bytes. + 5. **Release GPU** ā€” ``policy_generation.finish_generation()`` + frees KV cache and inference state so the trainer can use the + GPU immediately. + 6. **Capture metrics** ā€” ``policy_generation.get_logger_metrics()`` + collects generation stats (throughput, etc.) and returns them + to the driver in the result tuple. + + The driver receives ``(meta, driver_carry, rollout_metrics, + generation_logger_metrics)`` and uses ``driver_carry`` for its + own per-row compute (rewards, advantages, dynamic sampling). + + Args: + input_batch: Per-step prompt batch (already repeat-interleaved). + partition_id: TQ partition target. + group_size: Rollouts per original prompt. One uid is minted + per prompt; bulk keys are ``f"{uid}_g{i}"`` where ``i`` + ranges over the per-prompt expansion (group Ɨ rollout + turns). Train passes ``num_generations_per_prompt``; val + passes ``1``. + first_iter: True on the first DS iteration of a step; drives + ``policy_generation.snapshot_step_metrics()`` so per-step + metrics align with the legacy ``grpo.grpo_train`` path. + finish_generation: Call ``policy_generation.finish_generation()`` + at the tail. Default ``True`` matches the training step + (one rollout per step, release KV after). Validation sets + ``False`` so inference state survives across val batches; + the trainer owns the explicit ``finish_generation()`` call + at the end of the val pass. + task_to_env_override: Per-call task ā†’ env map. ``None`` uses + ``self.task_to_env`` (training envs supplied at construction). + Validation passes ``val_task_to_env`` here so val rollouts + run against the val env set without rebuilding the actor. + carry_keys: Names of per-row tensors to return in + ``driver_carry``. ``None`` returns every available key + (training uses this). Validation passes a slim list + (e.g. ``["total_reward"]``) to avoid wasting Ray transfer + on fields it doesn't consume. + + Returns: + ``(meta, driver_carry, rollout_metrics, generation_logger_metrics)`` + where ``driver_carry`` is a per-row dict of tensors the driver + uses for compute (rewards, masks, lengths, prompt_ids_for_adv, + ā€¦) ā€” stays on the driver, never crosses an actor boundary. + """ + # Lazy imports ā€” avoid pulling grpo into this module at load. + from nemo_rl.algorithms.grpo import ( + _extract_prompt_only_messages, + _should_use_async_rollouts, + _should_use_nemo_gym, + ) + from nemo_rl.algorithms.utils import get_gdpo_reward_component_keys + from nemo_rl.data.llm_message_utils import ( + MESSAGE_LOG_BULK_FIELDS, + add_loss_mask_to_message_log, + batched_message_log_to_flat_message, + decompose_message_log, + ) + + # Per-step generation-side metric hooks: snapshot once on the + # first DS iter so backends with per-step deltas have a stable + # anchor; clear accumulators before every rollout. Mirrors + # legacy ``grpo_train``. + if self.policy_generation is not None: + if first_iter and hasattr(self.policy_generation, "snapshot_step_metrics"): + self.policy_generation.snapshot_step_metrics() + self.policy_generation.clear_logger_metrics() + + cfg = self.master_config + task_to_env = ( + task_to_env_override + if task_to_env_override is not None + else self.task_to_env + ) + common = dict( + policy_generation=self.policy_generation, + input_batch=input_batch, + tokenizer=self.tokenizer, + task_to_env=task_to_env, + greedy=False, + ) + + # Rollout dispatch (mirrors grpo_sync.py:294-349). + if _should_use_nemo_gym(cfg): + r = run_async_nemo_gym_rollout( + **common, + max_seq_len=None, + max_rollout_turns=None, + generation_config=cfg.policy["generation"], + ) + final_batch, rollout_metrics = r.final_batch, r.rollout_metrics + else: + runner = ( + run_async_multi_turn_rollout + if _should_use_async_rollouts(cfg) + else run_multi_turn_rollout + ) + final_batch, rollout_metrics = runner( + **common, + max_seq_len=cfg.policy["max_total_sequence_length"], + max_rollout_turns=cfg.grpo["max_rollout_turns"], + ) + fb = final_batch.to("cpu") + del final_batch + + # Assistant-only loss mask (shared helper); seed missing + # generation_logprobs (e.g. when the env wraps assistant turns + # without a backing logprob, or for greedy/replay rollouts). + add_loss_mask_to_message_log(fb["message_log"]) + for ml in fb["message_log"]: + for msg in ml: + msg.setdefault( + "generation_logprobs", + torch.zeros_like(msg["token_ids"], dtype=torch.float32), + ) + + # Flatten message_log ā†’ bulk tensors + extract prompt-only ids. + pad = {"pad_value_dict": {"token_ids": self.tokenizer.pad_token_id}} + flat, input_lengths = batched_message_log_to_flat_message( + fb["message_log"], + **pad, + make_sequence_length_divisible_by=cfg.policy[ + "make_sequence_length_divisible_by" + ], + ) + prompt_flat, _ = batched_message_log_to_flat_message( + _extract_prompt_only_messages(fb["message_log"]), + **pad, + ) + + # TQ bulk payload ā€” DP_TRAIN_FIELDS + multimodal extras. + bulk_batch = BatchedDataDict[Any]( + { + "input_ids": flat["token_ids"], + "input_lengths": input_lengths, + "generation_logprobs": flat["generation_logprobs"], + "token_mask": flat["token_loss_mask"], + "sample_mask": fb["loss_multiplier"], + } + ) + for k, v in flat.get_multimodal_dict(as_tensors=False).items(): + if isinstance(v, torch.Tensor): + bulk_batch[k] = v + # ``content`` (raw assistant text per sample) ā€” rides TQ as a + # NonTensorStack so the driver can fetch it back at jsonl time + # (kv_first_write wraps it via NonTensorStack). + if "content" in flat: + bulk_batch["content"] = np.asarray(flat["content"], dtype=object) + + # Split `message_log` into per-field arrays instead of pickling + # the list-of-dicts-with-tensors per row. Consumer rebuilds + # `message_log` on read; external API stays the same. + decomposed = decompose_message_log(fb["message_log"]) + for k in MESSAGE_LOG_BULK_FIELDS: + bulk_batch[k] = decomposed[k] + + # Pass through remaining non-tensor fb fields as object arrays; + # `message_log` is excluded since its tensors live in the + # decomposed fields above (per-row pickle of dict-with-tensors + # would smuggle aliased views into the wire). + for k, v in fb.items(): + if isinstance(v, torch.Tensor) or k in bulk_batch or k == "message_log": + continue + bulk_batch[k] = ( + v + if isinstance(v, np.ndarray) and v.dtype == object + else np.asarray(v, dtype=object) + ) + + # Slice ā€” only what the driver can't derive from a TQ slice fetch + # (anything containing `message_log` or per-token data would + # force a fetch). Driver does scale_rewards / reward_shaping / + # overlong filtering / baseline-std on this slice. + truncated = fb["truncated"] + if not isinstance(truncated, torch.Tensor): + truncated = torch.tensor(truncated, dtype=torch.bool) + length = fb.get("length", input_lengths) + if not isinstance(length, torch.Tensor): + length = torch.tensor(length) + driver_carry = { + "total_reward": fb["total_reward"], + "loss_multiplier": fb["loss_multiplier"], + "truncated": truncated, + "length": length, + "input_lengths": input_lengths, + "prompt_ids_for_adv": prompt_flat["token_ids"], + # Computed by decompose_message_log above; feeds + # apply_reward_shaping on the driver without a TQ fetch. + "response_token_lengths": decomposed["response_token_lengths"], + } + # GDPO multi-reward components: scale_rewards iterates these + # keys driver-side and the GDPO advantage estimator reads them + # from ``adv_inputs``. Plumb them through ``driver_carry`` + # rather than forcing a separate TQ fetch. + for k in get_gdpo_reward_component_keys(fb): + driver_carry[k] = fb[k] + if carry_keys is not None: + for k in OPT_IN_CARRY_KEYS: + if k in carry_keys: + driver_carry[k] = decomposed[k] + missing = set(carry_keys) - driver_carry.keys() + if missing: + raise KeyError( + f"rollout_to_tq: carry_keys {sorted(missing)} not produced; " + f"valid keys: {sorted(driver_carry)}" + ) + driver_carry = {k: driver_carry[k] for k in carry_keys} + + n_samples = int(bulk_batch["sample_mask"].shape[0]) + input_size = int(input_batch.size) + if group_size <= 0 or input_size % group_size != 0: + raise ValueError( + f"input_batch.size={input_size} is not divisible by group_size={group_size}" + ) + n_prompts = input_size // group_size + if n_prompts == 0 or n_samples % n_prompts != 0: + raise ValueError( + f"bulk_batch has {n_samples} samples; not divisible by n_prompts={n_prompts}" + ) + n_per_prompt = n_samples // n_prompts + uids = [str(uuid.uuid4()) for _ in range(n_prompts)] + sample_ids = [f"{uid}_g{i}" for uid in uids for i in range(n_per_prompt)] + meta = kv_first_write( + bulk_batch, + sample_ids=sample_ids, + dp_client=self._dp_client, + partition_id=partition_id, + extra_info={"rollout_metrics": rollout_metrics}, + task_name=partition_id, + pad_to_multiple=int( + cfg.policy.get("make_sequence_length_divisible_by") or 1 + ), + ) + + if self.policy_generation is not None: + if finish_generation: + self.policy_generation.finish_generation() + gen_metrics = self.policy_generation.get_logger_metrics() + else: + gen_metrics = None + return meta, BatchedDataDict(driver_carry), rollout_metrics, gen_metrics + + def shutdown(self) -> None: + try: + self._dp_client.close() + except Exception: + pass diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index c3f7772c425..ea9b21d6a4a 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -367,6 +367,88 @@ def init_collective( # this function should co-work with vllm, so we should wait for all futures to complete outside return futures + # ā”€ā”€ DP-shard helpers ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + # DRY for Policy's logprob/train methods only. The data-plane sibling + # TQPolicy shards KVBatchMeta via ``shard_meta_for_dp``; the + # driver-on-data vs driver-on-meta split is by design. + def _shard_for_logprob( + self, + data: BatchedDataDict[Any], + ) -> tuple[list["SlicedDataDict"], Optional[list[int]]]: + """Shard inputs for ``get_logprobs`` / ``get_reference_policy_logprobs``. + + Mirrors the legacy shard block (lines 426-450 / 503-530). Returns + ``(sharded_data, unsorted_data_indices)`` where the second element + is the inverse permutation needed to undo seqpack/dynbatch reorder + (``None`` when neither is enabled). + """ + dp_size = self.sharding_annotations.get_axis_size("data_parallel") + if self.use_dynamic_batches: + self.dynamic_batching_args["max_tokens_per_microbatch"] = self.cfg[ + "dynamic_batching" + ]["logprob_mb_tokens"] + sharded_data, unsorted_data_indices = data.shard_by_batch_size( # type: ignore + dp_size, + batch_size=None, + dynamic_batching_args=self.dynamic_batching_args, + ) + elif self.use_sequence_packing: + self.sequence_packing_args["max_tokens_per_microbatch"] = self.cfg[ + "sequence_packing" + ]["logprob_mb_tokens"] + # we just shard into DP shards here as Sequence packing allows for CP. + sharded_data, unsorted_data_indices = data.shard_by_batch_size( + dp_size, + batch_size=None, + sequence_packing_args=self.sequence_packing_args, + ) + else: + sharded_data = data.shard_by_batch_size( # type: ignore + dp_size, + batch_size=None, + ) + unsorted_data_indices = None + return sharded_data, unsorted_data_indices + + def _shard_for_train( + self, + data: BatchedDataDict[Any], + batch_size: int, + ) -> list["SlicedDataDict"]: + """Shard inputs for ``train``. + + Mirrors the legacy shard block (lines 706-729). Note vs. + ``_shard_for_logprob``: uses ``train_mb_tokens`` (not + ``logprob_mb_tokens``), passes ``batch_size`` (not None), and + does not return ``unsorted_data_indices`` because train returns + scalar metrics (no per-row outputs to reorder). + """ + dp_size = self.sharding_annotations.get_axis_size("data_parallel") + if self.use_dynamic_batches: + self.dynamic_batching_args["max_tokens_per_microbatch"] = self.cfg[ + "dynamic_batching" + ]["train_mb_tokens"] + sharded_data, _ = data.shard_by_batch_size( + dp_size, + batch_size=batch_size, + dynamic_batching_args=self.dynamic_batching_args, + ) + elif self.use_sequence_packing: + self.sequence_packing_args["max_tokens_per_microbatch"] = self.cfg[ + "sequence_packing" + ]["train_mb_tokens"] + sharded_data, _ = data.shard_by_batch_size( + dp_size, + batch_size=batch_size, + sequence_packing_args=self.sequence_packing_args, + ) + else: + sharded_data = data.shard_by_batch_size( + dp_size, + batch_size=batch_size, + ) + return sharded_data + def get_logprobs( self, data: BatchedDataDict[GenerationDatumSpec], @@ -379,35 +461,8 @@ def get_logprobs( We use the convention that the logprob of the first token is 0 so that the sequence length is maintained. The logprob of input token i is specified at position i in the output logprobs tensor. """ - dp_size = self.sharding_annotations.get_axis_size("data_parallel") - sharded_data: list[SlicedDataDict] - unsorted_data_indices: list[int] - with timer.time("get_logprobs/shard_data") if timer else nullcontext(): - if self.use_dynamic_batches: - self.dynamic_batching_args["max_tokens_per_microbatch"] = self.cfg[ - "dynamic_batching" - ]["logprob_mb_tokens"] - sharded_data, unsorted_data_indices = data.shard_by_batch_size( # type: ignore - dp_size, - batch_size=None, - dynamic_batching_args=self.dynamic_batching_args, - ) - elif self.use_sequence_packing: - self.sequence_packing_args["max_tokens_per_microbatch"] = self.cfg[ - "sequence_packing" - ]["logprob_mb_tokens"] - # we just shard into DP shards here as Sequence packing allows for CP. - sharded_data, unsorted_data_indices = data.shard_by_batch_size( - dp_size, - batch_size=None, - sequence_packing_args=self.sequence_packing_args, - ) - else: - sharded_data = data.shard_by_batch_size( # type: ignore - dp_size, - batch_size=None, - ) + sharded_data, unsorted_data_indices = self._shard_for_logprob(data) with ( timer.time("get_logprobs/submit_logprob_futures") @@ -435,7 +490,7 @@ def get_logprobs( # dynamic batching sorts the inputs by sequence length to improve load balancing, # so change it back here - if self.use_dynamic_batches or self.use_sequence_packing: + if unsorted_data_indices is not None: logprobs.reorder_data(unsorted_data_indices) return logprobs @@ -450,37 +505,12 @@ def get_reference_policy_logprobs( Returns: Identical to get_logprobs. """ - dp_size = self.sharding_annotations.get_axis_size("data_parallel") - sharded_data: list[SlicedDataDict] - unsorted_data_indices: list[int] with ( timer.time("get_reference_policy_logprobs/shard_data") if timer else nullcontext() ): - if self.use_dynamic_batches: - self.dynamic_batching_args["max_tokens_per_microbatch"] = self.cfg[ - "dynamic_batching" - ]["logprob_mb_tokens"] - sharded_data, unsorted_data_indices = data.shard_by_batch_size( # type: ignore - dp_size, - batch_size=None, - dynamic_batching_args=self.dynamic_batching_args, - ) - elif self.use_sequence_packing: - self.sequence_packing_args["max_tokens_per_microbatch"] = self.cfg[ - "sequence_packing" - ]["logprob_mb_tokens"] - sharded_data, unsorted_data_indices = data.shard_by_batch_size( - dp_size, - batch_size=None, - sequence_packing_args=self.sequence_packing_args, - ) - else: - sharded_data = data.shard_by_batch_size( # type: ignore - dp_size, - batch_size=None, - ) + sharded_data, unsorted_data_indices = self._shard_for_logprob(data) with ( timer.time( @@ -513,7 +543,7 @@ def get_reference_policy_logprobs( # dynamic batching sorts the inputs by sequence length to improve load balancing, # so change it back here - if self.use_dynamic_batches or self.use_sequence_packing: + if unsorted_data_indices is not None: logprobs.reorder_data(unsorted_data_indices) return logprobs @@ -526,34 +556,8 @@ def get_topk_logits( timer: Optional[Timer] = None, ) -> BatchedDataDict[TopkLogitsOutputSpec]: """Dispatch get_topk_logits to workers (no CP/packed support initially).""" - dp_size = self.sharding_annotations.get_axis_size("data_parallel") - sharded_data: list[SlicedDataDict] - unsorted_data_indices: list[int] with timer.time("get_topk_logits/shard_data") if timer else nullcontext(): - if self.use_dynamic_batches: - self.dynamic_batching_args["max_tokens_per_microbatch"] = self.cfg[ - "dynamic_batching" - ]["logprob_mb_tokens"] - sharded_data, unsorted_data_indices = data.shard_by_batch_size( # type: ignore - dp_size, - batch_size=None, - dynamic_batching_args=self.dynamic_batching_args, - ) - elif self.use_sequence_packing: - self.sequence_packing_args["max_tokens_per_microbatch"] = self.cfg[ - "sequence_packing" - ]["logprob_mb_tokens"] - # we just shard into DP shards here as Sequence packing allows for CP. - sharded_data, unsorted_data_indices = data.shard_by_batch_size( - dp_size, - batch_size=None, - sequence_packing_args=self.sequence_packing_args, - ) - else: - sharded_data = data.shard_by_batch_size( # type: ignore - dp_size, - batch_size=None, - ) + sharded_data, unsorted_data_indices = self._shard_for_logprob(data) with ( timer.time("get_topk_logits/submit_topk_logits_futures") @@ -586,7 +590,7 @@ def get_topk_logits( stacked["topk_logits"] = torch.cat(all_topk_logits, dim=0) stacked["topk_indices"] = torch.cat(all_topk_indices, dim=0) - if self.use_dynamic_batches or self.use_sequence_packing: + if unsorted_data_indices is not None: stacked.reorder_data(unsorted_data_indices) return stacked @@ -604,31 +608,8 @@ def train( batch_size = gbs or self.cfg["train_global_batch_size"] micro_batch_size = mbs or self.cfg["train_micro_batch_size"] # Shard and replicate the batch - dp_size = self.sharding_annotations.get_axis_size("data_parallel") with timer.time("policy_training/sharding_data") if timer else nullcontext(): - if self.use_dynamic_batches: - self.dynamic_batching_args["max_tokens_per_microbatch"] = self.cfg[ - "dynamic_batching" - ]["train_mb_tokens"] - sharded_data, _ = data.shard_by_batch_size( - dp_size, - batch_size=batch_size, - dynamic_batching_args=self.dynamic_batching_args, - ) - elif self.use_sequence_packing: - self.sequence_packing_args["max_tokens_per_microbatch"] = self.cfg[ - "sequence_packing" - ]["train_mb_tokens"] - sharded_data, _ = data.shard_by_batch_size( - dp_size, - batch_size=batch_size, - sequence_packing_args=self.sequence_packing_args, - ) - else: - sharded_data = data.shard_by_batch_size( - dp_size, - batch_size=batch_size, - ) + sharded_data = self._shard_for_train(data, batch_size) if self.flops_tracker is not None: self.flops_tracker.reset() diff --git a/nemo_rl/models/policy/tq_policy.py b/nemo_rl/models/policy/tq_policy.py new file mode 100644 index 00000000000..1179bd8a1fc --- /dev/null +++ b/nemo_rl/models/policy/tq_policy.py @@ -0,0 +1,440 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""TQ-mediated Policy: meta-driven 1-hop counterpart to ``Policy``. + +Exposes ``train_from_meta`` / ``get_logprobs_from_meta`` / +``get_reference_policy_logprobs_from_meta`` ā€” same return shapes as +``Policy.{train, get_logprobs, get_reference_policy_logprobs}`` but +accepting a ``KVBatchMeta`` instead of a ``BatchedDataDict``. The meta +names per-sample TQ keys minted once at rollout +(:class:`nemo_rl.experience.sync_rollout_actor.SyncRolloutActor`); each +dispatch slices the key list per DP rank via +:func:`nemo_rl.data_plane.preshard.shard_meta_for_dp` (no re-fan-out, +no key minting). Workers fetch their slice from TQ via +``self._fetch(meta)`` and write deltas back via +``self._write_back_result_field(...)``. See +``nemo_rl/data_plane/README.md`` for the full design. +""" + +from __future__ import annotations + +import warnings +from collections import defaultdict +from contextlib import nullcontext +from dataclasses import replace +from typing import Any, Optional + +import ray + +from nemo_rl.algorithms.loss.interfaces import LossFunction +from nemo_rl.data_plane import KVBatchMeta, build_data_plane_client +from nemo_rl.data_plane.column_io import read_columns, round_up, write_columns +from nemo_rl.data_plane.preshard import shard_meta_for_dp +from nemo_rl.data_plane.schema import ( + DP_TRAIN_FIELDS, + GLOBAL_FORWARD_PAD_SEQLEN, + LP_SEED_FIELDS, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.models.policy.lm_policy import Policy +from nemo_rl.utils.flops_tracker import get_theoretical_tflops +from nemo_rl.utils.timer import Timer + +# ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ +# Per-stage aggregators that assemble per-rank worker results into the +# shape each Policy method returns. Used by the TQ-mediated overrides +# below; kept out of ``lm_policy.Policy`` since the legacy in-memory +# path doesn't fan out per-rank and never calls these. +# ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def _aggregate_train_results(results: list[dict[str, Any]]) -> dict[str, Any]: + out: dict[str, Any] = { + "loss": results[0]["global_loss"], + "grad_norm": results[0]["grad_norm"], + } + if "moe_metrics" in results[0]: + out["moe_metrics"] = results[0]["moe_metrics"] + all_mb_metrics: dict[str, list[Any]] = defaultdict(list) + for r in results: + for k, v in r["all_mb_metrics"].items(): + all_mb_metrics[k].extend(v) + out["all_mb_metrics"] = dict(all_mb_metrics) + return out + + +# Logprob results land in TQ directly via the worker-side +# ``_write_back_result_field`` leader path; the per-rank Ray return is +# always None (see :meth:`TQWorkerMixin.get_logprobs_presharded`). The +# dispatcher only waits for completion ā€” no aggregation needed. + + +class TQPolicy(Policy): + """TQ-mediated counterpart to :class:`Policy`. + + Constructor accepts an additional ``dp_cfg`` (the + ``master_config["data_plane"]`` dict). Bootstraps the controller on + the driver and forwards ``setup_data_plane(dp_cfg)`` to every worker + so they can attach as clients (``bootstrap=False``). + + The partition lifecycle (``register_partition`` / ``clear_samples``) is + the trainer's responsibility ā€” this class assumes the partition + named ``self.tq_partition_id`` (default ``"train"``) is open with a + schema covering ``DP_TRAIN_FIELDS`` (the bulk schema written by the + rollout actor at first put + driver-/worker-written deltas). + """ + + def __init__( + self, + *args: Any, + dp_cfg: dict[str, Any], + tq_partition_id: str = "train", + **kwargs: Any, + ) -> None: + super().__init__(*args, **kwargs) + # Validate the topology the data plane fan-out (`shard_meta_for_dp`) + # depends on. Failing here surfaces a clear error at policy + # construction; the same condition is re-checked inside + # `shard_meta_for_dp` as a defensive invariant. + dp_world = self.sharding_annotations.get_axis_size("data_parallel") + if dp_world <= 0: + raise ValueError( + f"TQPolicy requires data_parallel axis size > 0, got {dp_world}. " + f"Check cluster config (gpus_per_node * num_nodes) vs. " + f"TP/PP/CP/EP sizes." + ) + self.dp_cfg = dp_cfg + self.dp_client = build_data_plane_client(dp_cfg, bootstrap=True) + self.tq_partition_id = tq_partition_id + + # Forward to workers (replaces ``Policy.setup_data_plane`` call + # site in the trainer ā€” TQPolicy bundles bootstrap + worker + # attach into construction so the trainer just instantiates + # ``TQPolicy(...)`` and is done). + ray.get( + self.worker_group.run_all_workers_single_data( + "setup_data_plane", cfg=dp_cfg + ) + ) + + # ā”€ā”€ lifecycle ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + def shutdown(self) -> bool: # type: ignore[override] + """Close the TQ client before shutting down the worker group.""" + try: + self.dp_client.close() + except Exception as e: + warnings.warn(f"Error closing data-plane client: {e}") + return super().shutdown() + + def prepare_step( + self, + num_samples: int, + group_size: Optional[int] = None, + ) -> None: + """Register the per-step TQ partition. + + Sync trainers call this at the start of each step. The static + partition id ``"train"`` is cleared and reused across steps. The + schema is the union of all consumer fields ā€” producers write + only the subset they have, consumers fetch via ``select_fields``. + + Args: + num_samples: Expected total samples this step. + group_size: GRPO group size for balanced sampling; ``None`` disables grouping. + """ + self.dp_client.register_partition( + partition_id=self.tq_partition_id, + fields=list(DP_TRAIN_FIELDS), + num_samples=num_samples, + consumer_tasks=["prev_lp", "ref_lp", "train"], + grpo_group_size=group_size, + ) + + def prepare_val_partition( + self, num_samples: int, *, partition_id: str = "val" + ) -> None: + """Register a per-batch val partition (single consumer, no GRPO grouping). + + Sync val trainers call this at the start of each val batch. + Distinct from :meth:`prepare_step` because val has its own + partition id and a single consumer task. + """ + self.dp_client.register_partition( + partition_id=partition_id, + fields=list(DP_TRAIN_FIELDS), + num_samples=num_samples, + consumer_tasks=[partition_id], + grpo_group_size=None, + ) + + def discard_samples(self, sample_ids: list[str], partition_id: str) -> None: + """Drop a set of uids from TQ. + + Used both for step-end teardown (via :meth:`finish_step`) and + mid-step filtering (e.g. dynamic sampling). + """ + self.dp_client.clear_samples(sample_ids=sample_ids, partition_id=partition_id) + + def finish_step(self, meta: KVBatchMeta) -> None: + """Drop this step's bulk from TQ. Mirror of :meth:`prepare_step`.""" + self.discard_samples(meta.sample_ids, meta.partition_id) + + def _stamp_pad_seqlen(self, meta: KVBatchMeta) -> None: + """Mint ``GLOBAL_FORWARD_PAD_SEQLEN`` onto ``meta.extra_info`` (idempotent). + + Cross-DP forward pad target. Preshard shards inherit it via + ``dict(meta.extra_info)`` propagation. + """ + if not meta.sequence_lengths: + return + if GLOBAL_FORWARD_PAD_SEQLEN in meta.extra_info: + return + _, dba = self._packing_args("train_mb_tokens") + seq_round = int(dba["sequence_length_round"]) if dba is not None else 1 + pad_mult = int(meta.extra_info.get("pad_to_multiple", 1)) + meta.extra_info[GLOBAL_FORWARD_PAD_SEQLEN] = round_up( + max(meta.sequence_lengths), max(pad_mult, seq_round) + ) + + def read_from_dataplane( + self, + meta: KVBatchMeta, + *, + select_fields: list[str], + pad_value_dict: Optional[dict[str, Any]] = None, + ) -> BatchedDataDict[Any]: + """Fetch + materialize columns from the data plane (TQ). + + ``read_columns`` pads to ``meta.extra_info[GLOBAL_FORWARD_PAD_SEQLEN]`` + ā€” the same value workers pad to in their forward pass. Driver + and workers thus return columns at one identical seq dim, with + no driver-side knowledge of ``sequence_length_round``. + """ + self._stamp_pad_seqlen(meta) + return read_columns( + self.dp_client, + meta, + select_fields=select_fields, + pad_value_dict=pad_value_dict, + ) + + def write_to_dataplane(self, meta: KVBatchMeta, fields: dict[str, Any]) -> None: + """Write driver-computed columns to the data plane (TQ).""" + write_columns(self.dp_client, meta, fields=fields) + + # ā”€ā”€ 1-hop entrypoints (KVBatchMeta in, no re-fan-out) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + def _packing_args( + self, + mb_tokens_key: str, + ) -> tuple[Optional[dict[str, Any]], Optional[dict[str, Any]]]: + """Resolve (sequence_packing_args, dynamic_batching_args) for a given stage. + + The stage is identified by ``mb_tokens_key`` (``"logprob_mb_tokens"`` or + ``"train_mb_tokens"``). + """ + if getattr(self, "use_dynamic_batches", False): + args = dict(self.dynamic_batching_args) + args["max_tokens_per_microbatch"] = self.cfg["dynamic_batching"][ + mb_tokens_key + ] + return None, args + if getattr(self, "use_sequence_packing", False): + args = dict(self.sequence_packing_args) + args["max_tokens_per_microbatch"] = self.cfg["sequence_packing"][ + mb_tokens_key + ] + return args, None + return None, None + + def _logprob_dispatch( + self, + meta: KVBatchMeta, + *, + task_name: str, + worker_method: str, + timer_prefix: str, + timer: Optional[Timer], + common_kwargs: dict[str, Any], + ) -> None: + """Shared body of get_logprobs_from_meta / get_reference_policy_logprobs_from_meta. + + Logprob workers need only LP_SEED_FIELDS ā€” narrow the meta's + field list so ``_fetch`` doesn't pull rollout-only payload (e.g. + multimodal). The same shape is used for both prev_lp and ref_lp. + Workers compute the per-token tensor and commit it to TQ via the + leader-rank ``_write_back_result_field``; the Ray return is + always None, so this dispatcher just waits for completion. + """ + self._stamp_pad_seqlen(meta) + spa, dba = self._packing_args("logprob_mb_tokens") + lp_meta = replace(meta, fields=list(LP_SEED_FIELDS), task_name=task_name) + with timer.time(f"{timer_prefix}/shard_meta") if timer else nullcontext(): + metas, _ = shard_meta_for_dp( + lp_meta, + dp_world=self.sharding_annotations.get_axis_size("data_parallel"), + batch_size=None, + sequence_packing_args=spa, + dynamic_batching_args=dba, + ) + with timer.time(f"{timer_prefix}/submit_futures") if timer else nullcontext(): + futures = self.worker_group.run_all_workers_sharded_data( + worker_method, + meta=metas, + in_sharded_axes=["data_parallel"], + replicate_on_axes=[ + "context_parallel", + "tensor_parallel", + "pipeline_parallel", + ], + output_is_replicated=[ + "context_parallel", + "tensor_parallel", + "pipeline_parallel", + ], + common_kwargs=common_kwargs, + ) + # Wait for completion; per-rank returns are None. + self.worker_group.get_all_worker_results(futures) + + def get_logprobs_from_meta( + self, + meta: KVBatchMeta, + micro_batch_size: Optional[int] = None, + timer: Optional[Timer] = None, + ) -> None: + self._logprob_dispatch( + meta, + task_name="prev_lp", + worker_method="get_logprobs_presharded", + timer_prefix="get_logprobs", + timer=timer, + common_kwargs={"micro_batch_size": micro_batch_size}, + ) + + def get_reference_policy_logprobs_from_meta( + self, + meta: KVBatchMeta, + micro_batch_size: Optional[int] = None, + timer: Optional[Timer] = None, + ) -> None: + self._logprob_dispatch( + meta, + task_name="ref_lp", + worker_method="get_reference_policy_logprobs_presharded", + timer_prefix="get_reference_policy_logprobs", + timer=timer, + common_kwargs={"micro_batch_size": micro_batch_size}, + ) + + def train_from_meta( + self, + meta: KVBatchMeta, + loss_fn: LossFunction, + eval_mode: bool = False, + gbs: Optional[int] = None, + mbs: Optional[int] = None, + timer: Optional[Timer] = None, + ) -> dict[str, Any]: + """1-hop counterpart to :meth:`train`. + + ``meta`` names per-sample keys; columns written by the rollout + actor + worker logprob deltas + driver-side advantage delta have + all landed under the same keys at this point. Workers fetch the + union via ``train_presharded`` ā†’ ``self._fetch(meta)``. No + partition drain here ā€” sync 1-hop's trainer calls ``clear_samples`` + once at end of step. + + Args: + meta: Full-step ``KVBatchMeta`` (consumed by all DP ranks). + gbs: Global batch size; defaults to ``cfg["train_global_batch_size"]``. + mbs: Micro batch size; defaults to ``cfg["train_micro_batch_size"]``. + timer: Optional timer for nested ``policy_training/*`` measurements. + + Returns: + Aggregated training-step output dict. + """ + batch_size = gbs or self.cfg["train_global_batch_size"] + micro_batch_size = mbs or self.cfg["train_micro_batch_size"] + + self._stamp_pad_seqlen(meta) + spa, dba = self._packing_args("train_mb_tokens") + # Train workers fetch the full DP_TRAIN_FIELDS schema (rollout + + # logprob deltas + advantages + sample_mask). Caller is responsible + # for ensuring those columns have been written to TQ before this + # call (workers + driver delta-writes). + train_meta = replace( + meta, + fields=list(DP_TRAIN_FIELDS), + task_name="train", + ) + with timer.time("policy_training/shard_meta") if timer else nullcontext(): + dp_metas, _ = shard_meta_for_dp( + train_meta, + dp_world=self.sharding_annotations.get_axis_size("data_parallel"), + batch_size=batch_size, + sequence_packing_args=spa, + dynamic_batching_args=dba, + ) + + if self.flops_tracker is not None: + self.flops_tracker.reset() + for m in dp_metas: + self.flops_tracker.track_batch(list(m.sequence_lengths or [])) + + with ( + timer.time("policy_training/submit_training_futures") + if timer + else nullcontext() + ): + futures = self.worker_group.run_all_workers_sharded_data( + "train_presharded", + meta=dp_metas, + in_sharded_axes=["data_parallel"], + replicate_on_axes=[ + "context_parallel", + "tensor_parallel", + "pipeline_parallel", + ], + output_is_replicated=[ + "context_parallel", + "tensor_parallel", + "pipeline_parallel", + ], + common_kwargs={ + "loss_fn": loss_fn, + "eval_mode": eval_mode, + "gbs": batch_size, + "mbs": micro_batch_size, + }, + ) + results = self.worker_group.get_all_worker_results(futures) + aggregated_results = _aggregate_train_results(results) + + if self.flops_tracker is not None: + aggregated_results["total_flops"] = self.flops_tracker.total_flops + aggregated_results["num_ranks"] = self.worker_group.cluster.world_size() + gpus_per_worker = self.worker_group.cluster.world_size() / max( + len(results), 1 + ) + try: + aggregated_results["theoretical_tflops"] = gpus_per_worker * sum( + get_theoretical_tflops(r["gpu_name"], r["model_dtype"]) + for r in results + ) + except Exception as e: + warnings.warn(f"Error getting theoretical flops: {e}") + + return aggregated_results diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker.py b/nemo_rl/models/policy/workers/dtensor_policy_worker.py index 022335f7d02..bb1b9e52f58 100644 --- a/nemo_rl/models/policy/workers/dtensor_policy_worker.py +++ b/nemo_rl/models/policy/workers/dtensor_policy_worker.py @@ -58,6 +58,7 @@ from nemo_rl.algorithms.loss import SequencePackingLossWrapper, prepare_loss_input from nemo_rl.algorithms.loss.interfaces import LossFunction, LossType from nemo_rl.algorithms.utils import mask_out_neg_inf_logprobs +from nemo_rl.data_plane.worker_mixin import TQWorkerMixin from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.model_utils import ( allgather_cp_sharded_tensor, @@ -164,7 +165,9 @@ def get_cpu_state_dict( # Classes with @ray.remote can't be inherited from, so we split the implementation out. # This is useful when using worker extension classes. -class DTensorPolicyWorkerImpl(AbstractPolicyWorker, ColocatablePolicyInterface): +class DTensorPolicyWorkerImpl( + TQWorkerMixin, AbstractPolicyWorker, ColocatablePolicyInterface +): def __repr__(self) -> str: """Customizes the actor's prefix in the Ray logs. @@ -175,6 +178,16 @@ def __repr__(self) -> str: else: return f"{self.__class__.__qualname__}" + def _get_replica_group(self) -> Optional[Any]: + """Replica group = flattened (cp, tp) sub-mesh, for NCCL broadcast in ``_fetch``.""" + return self.device_mesh[("cp", "tp")]._flatten().get_group() + + def _local_coords(self) -> dict[str, int]: + return { + "tensor_parallel": self.device_mesh["tp"].get_local_rank(), + "context_parallel": self.device_mesh["cp"].get_local_rank(), + } + def __init__( self, config: PolicyConfig, diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py b/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py index 2fa8a8e6043..27803e126bc 100644 --- a/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py +++ b/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py @@ -37,6 +37,7 @@ from nemo_rl.algorithms.logits_sampling_utils import TrainingSamplingParams from nemo_rl.algorithms.loss.interfaces import LossFunction +from nemo_rl.data_plane.worker_mixin import TQWorkerMixin from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.models.automodel.checkpoint import AutomodelCheckpointManager from nemo_rl.models.automodel.data import ( @@ -190,7 +191,9 @@ def get_train_context( # Classes with @ray.remote can't be inherited from, so we split the implementation out. # This is useful when using worker extension classes. -class DTensorPolicyWorkerV2Impl(AbstractPolicyWorker, ColocatablePolicyInterface): +class DTensorPolicyWorkerV2Impl( + TQWorkerMixin, AbstractPolicyWorker, ColocatablePolicyInterface +): def __repr__(self) -> str: """Customizes the actor's prefix in the Ray logs. @@ -201,6 +204,16 @@ def __repr__(self) -> str: else: return f"{self.__class__.__qualname__}" + def _get_replica_group(self) -> Optional[Any]: + """Replica group = flattened (cp, tp) sub-mesh ā€” see V1 worker.""" + return self.device_mesh[("cp", "tp")]._flatten().get_group() + + def _local_coords(self) -> dict[str, int]: + return { + "tensor_parallel": self.device_mesh["tp"].get_local_rank(), + "context_parallel": self.device_mesh["cp"].get_local_rank(), + } + def __init__( self, config: PolicyConfig, diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index c8a131cdc04..9ede82d0715 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -46,6 +46,7 @@ from nemo_rl.algorithms.logits_sampling_utils import TrainingSamplingParams from nemo_rl.algorithms.loss.interfaces import LossFunction +from nemo_rl.data_plane.worker_mixin import TQWorkerMixin from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.named_sharding import NamedSharding from nemo_rl.models.generation.interfaces import ( @@ -97,7 +98,9 @@ # Classes with @ray.remote can't be inherited from, so we split the implementation out. # This is useful when using worker extension classes. -class MegatronPolicyWorkerImpl(AbstractPolicyWorker, ColocatablePolicyInterface): +class MegatronPolicyWorkerImpl( + TQWorkerMixin, AbstractPolicyWorker, ColocatablePolicyInterface +): def __repr__(self): """Customizes the actor's prefix in the Ray logs. @@ -108,6 +111,66 @@ def __repr__(self): else: return f"{self.__class__.__qualname__}" + def _local_coords(self) -> dict[str, int]: + if not torch.distributed.is_initialized(): + return {} + return { + "tensor_parallel": parallel_state.get_tensor_model_parallel_rank(), + "context_parallel": parallel_state.get_context_parallel_rank(), + "pipeline_parallel": parallel_state.get_pipeline_model_parallel_rank(), + } + + def _get_replica_group(self) -> Optional[Any]: + """Replica group = TP Ɨ CP Ɨ PP siblings within this DP rank. + + Always returns the real group so :meth:`_is_replica_leader` (used + by both fetch and write-back) gives the correct single-writer + answer even at CP=1 ā€” gating on CP=1 here is what produced the + ``-601 ILLEGAL_CLIENT`` duplicate-write bug. The fetch-path + broadcast-vs-independent perf choice lives inside ``_fetch`` + keyed on ``replica_group.size()``. + + mcore exposes per-axis groups (``get_tensor_model_parallel_group``, + ``get_context_parallel_group``, ``get_pipeline_model_parallel_group``) + but no single combined group. We build the combined NCCL group + once on first call by enumerating coordinates that share this + rank's ``dp_rank``. + """ + if not torch.distributed.is_initialized(): + return None + cached = getattr(self, "_replica_group_cache", "uninit") + if cached != "uninit": + return cached + + world_size = torch.distributed.get_world_size() + my_dp_rank = parallel_state.get_data_parallel_rank() + # Collect global ranks that share this DP rank ā€” they form the + # replica group. Done collectively so every rank ends up with + # the same ranks list and can pass it to new_group(). + my_replica_ranks_t = torch.full( + (world_size,), + -1, + dtype=torch.long, + device="cuda", + ) + my_replica_ranks_t[torch.distributed.get_rank()] = my_dp_rank + torch.distributed.all_reduce( + my_replica_ranks_t, op=torch.distributed.ReduceOp.MAX + ) + all_dp_ranks = my_replica_ranks_t.tolist() + + # Every (dp_rank ā†’ ranks) bucket must call new_group on its own + # ranks list, but new_group itself must be called collectively + # across the full world. Sort by dp_rank to keep call order + # consistent across processes. + groups: dict[int, Any] = {} + for dp in sorted(set(all_dp_ranks)): + ranks = [r for r, d in enumerate(all_dp_ranks) if d == dp] + grp = torch.distributed.new_group(ranks=ranks, backend="nccl") + groups[dp] = grp + self._replica_group_cache = groups[my_dp_rank] + return self._replica_group_cache + def __init__( self, config: PolicyConfig, diff --git a/nemo_rl/utils/venvs.py b/nemo_rl/utils/venvs.py index 72e5ca39c46..c5fee47aa07 100644 --- a/nemo_rl/utils/venvs.py +++ b/nemo_rl/utils/venvs.py @@ -191,3 +191,35 @@ def create_local_venv_on_each_node(py_executable: str, venv_name: str): ray.util.remove_placement_group(pg) # Return mapping from node IP to venv python path return paths[0] + + +def make_actor_runtime_env(actor_class_fqn: str) -> dict: + """Build a Ray ``runtime_env`` for one of our registered actors. + + Resolves the actor's tier-specific py_executable via the registry, + materializes a per-node venv when uv-managed, and packages it with + ``VIRTUAL_ENV`` / ``UV_PROJECT_ENVIRONMENT`` env vars so workers see + the same interpreter as the driver. + + Used by ReplayBuffer, AsyncTrajectoryCollector, and + SyncRolloutActor ā€” three actors that need the VLLM tier's + venv on every node. + """ + # Local import ā€” venvs.py is dep-light; the registry imports + # PY_EXECUTABLES which transitively pulls heavier deps. + from nemo_rl.distributed.ray_actor_environment_registry import ( + get_actor_python_env, + ) + + py_exec = get_actor_python_env(actor_class_fqn) + if py_exec.startswith("uv"): + py_exec = create_local_venv_on_each_node(py_exec, actor_class_fqn) + venv = os.path.dirname(os.path.dirname(py_exec)) # strip bin/python + return { + "py_executable": py_exec, + "env_vars": { + **os.environ, + "VIRTUAL_ENV": venv, + "UV_PROJECT_ENVIRONMENT": venv, + }, + } diff --git a/pyproject.toml b/pyproject.toml index 59ad05b9a17..c3f356fd8a1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,41 @@ dependencies = [ "cuda-bindings; sys_platform != 'darwin'", # for non-colocated refit "pybase64", # for sglang refit "nvidia-cudnn-cu13==9.20.0.48; sys_platform != 'darwin'", # for transformer-engine no build isolation + # tilelang ā€” replacement Triton kernel mamba-ssm requires when + # Triton >= 3.4.0 on Hopper, see github.com/state-spaces/mamba#640. + # Without this, qwen3.5 / nano-v3 / moonlight megatron recipes + # crash at first gated-chunk backward with a RuntimeError pointing + # at this exact pip install. Linux x86_64 only ā€” mamba-ssm itself + # is gated to that pair. + "tilelang; sys_platform == 'linux' and platform_machine == 'x86_64'", + # Data-plane stack ā€” promoted to base so worker venvs (built by + # nemo_rl.utils.venvs.create_local_venv via bare `uv sync`, no extras) + # automatically include them. Removes the need for a `[data-plane]` + # extra and the corresponding plumbing in the per-worker venv builder. + "tensordict", + # Pinned to b266d39 (post-0.1.6, pre-0.1.7) for PR #77's MooncakeStore + # refactor: `clear` switched from unanchored `remove_by_regex` to + # exact-key `batch_remove`, which fixes a collateral-key-deletion bug + # that breaks DAPO + mooncake_cpu. Bump to the 0.1.7 tag when released. + "TransferQueue @ git+https://github.com/Ascend/TransferQueue.git@b266d39", + # Backs data_plane.backend="mooncake_cpu". Default backend is "simple" + # (in-process), but the mooncake_cpu path needs the `mooncake_master` + # binary that ships in this wheel at /mooncake/. Bundled + # with TQ rather than gated behind an extra so worker venvs (built + # without extras) can be flipped to mooncake_cpu via config alone. + # PyPI's base `mooncake-transfer-engine` is cu12-only (links + # libcudart.so.12), which breaks on cu13 containers. Upstream now also + # publishes a cu13 variant as a separate distribution name + # `mooncake-transfer-engine-cuda13` (same `mooncake/` import namespace, + # store.so linked against libcudart.so.13). Resolve from PyPI rather + # than the GitHub release URL ā€” the wheel is byte-identical (verified + # sha256), and PyPI's CDN is far more reliable than github releases + # from compute nodes. + # Upstream publishes both x86_64 and aarch64 wheels (see uv.lock). CI's + # build-container runner is aarch64 (uv reports aarch64-unknown-linux-gnu), + # so the marker must include aarch64 ā€” otherwise mooncake is silently + # excluded from the resolution during the Docker build. + "mooncake-transfer-engine-cuda13==0.3.10.post2 ; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", ] [project.optional-dependencies] @@ -304,6 +339,11 @@ override-dependencies = [ "outlines>=0.2.0", # Upgrade pytest to 9.0.3 "pytest>=9.0.3", + # TransferQueue (data-plane extra) pins numpy<2.0.0; megatron-core needs + # numpy>=2.1.0 via onnx ā†’ ml-dtypes. Override globally so the data-plane + # extra composes with mcore/automodel without version-mirroring TQ's + # requirements.txt. Forward-compatible across TQ minor bumps. + "numpy>=2.1.0", ] # CVE fixes diff --git a/pyrefly.toml b/pyrefly.toml index d79920b67eb..4d14b6d46b5 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -91,6 +91,18 @@ project-includes = [ "nemo_rl/data/multimodal_utils.py", "nemo_rl/data/packing/__init__.py", "nemo_rl/data/processors.py", + "nemo_rl/data_plane/__init__.py", + "nemo_rl/data_plane/adapters/__init__.py", + "nemo_rl/data_plane/adapters/noop.py", + "nemo_rl/data_plane/adapters/transfer_queue.py", + "nemo_rl/data_plane/codec.py", + "nemo_rl/data_plane/column_io.py", + "nemo_rl/data_plane/factory.py", + "nemo_rl/data_plane/interfaces.py", + "nemo_rl/data_plane/observability.py", + "nemo_rl/data_plane/preshard.py", + "nemo_rl/data_plane/schema.py", + "nemo_rl/data_plane/worker_mixin.py", "nemo_rl/distributed/__init__.py", "nemo_rl/distributed/collectives.py", "nemo_rl/distributed/named_sharding.py", diff --git a/tests/functional/L1_Functional_Tests_GPU.sh b/tests/functional/L1_Functional_Tests_GPU.sh index 57bc33bffa6..af5ebbb7d45 100644 --- a/tests/functional/L1_Functional_Tests_GPU.sh +++ b/tests/functional/L1_Functional_Tests_GPU.sh @@ -51,6 +51,8 @@ run_test fast uv run --no-sync bash ./tests/functional/eval_audio.sh run_test fast uv run --no-sync bash ./tests/functional/gdpo.sh run_test fast uv run --no-sync bash ./tests/functional/gdpo_async_grpo.sh run_test fast uv run --no-sync bash ./tests/functional/grpo.sh +run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_simple.sh +run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_mooncake.sh run_test fast uv run --no-sync bash ./tests/functional/grpo_async_gym.sh run_test uv run --no-sync bash ./tests/functional/grpo_automodel_lora.sh run_test uv run --no-sync bash ./tests/functional/grpo_automodel_lora_async.sh diff --git a/tests/functional/grpo_dp_mooncake.sh b/tests/functional/grpo_dp_mooncake.sh new file mode 100755 index 00000000000..b646f6c75a8 --- /dev/null +++ b/tests/functional/grpo_dp_mooncake.sh @@ -0,0 +1,51 @@ +#!/bin/bash +# Lightweight e2e for grpo_sync.py ā€” TQ pipeline with the mooncake_cpu +# backend. Same shape as tests/functional/grpo.sh (Qwen3-0.6B, 2 GPUs, +# 2 steps); exercises the real Mooncake transfer engine on the CPU path. + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) +# Mark the current repo as safe, since wandb fetches metadata about the repo +git config --global --add safe.directory $PROJECT_ROOT + +set -eou pipefail + +EXP_NAME=$(basename $0 .sh) +EXP_DIR=$SCRIPT_DIR/$EXP_NAME +LOG_DIR=$EXP_DIR/logs +JSON_METRICS=$EXP_DIR/metrics.json +RUN_LOG=$EXP_DIR/run.log +export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} + +rm -rf $EXP_DIR $LOG_DIR +mkdir -p $EXP_DIR $LOG_DIR + +cd $PROJECT_ROOT +uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ + $PROJECT_ROOT/examples/run_grpo.py \ + policy.model_name=Qwen/Qwen3-0.6B \ + grpo.num_prompts_per_step=2 \ + grpo.num_generations_per_prompt=4 \ + policy.train_global_batch_size=4 \ + policy.train_micro_batch_size=1 \ + cluster.gpus_per_node=2 \ + grpo.max_num_steps=2 \ + logger.tensorboard_enabled=true \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=false \ + logger.monitor_gpus=true \ + checkpointing.enabled=false \ + data_plane.enabled=true \ + data_plane.impl=transfer_queue \ + data_plane.backend=mooncake_cpu \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +uv run tests/check_metrics.py $JSON_METRICS \ + 'max(data["train/gen_kl_error"]) < 0.002' \ + 'min(data["train/probs_ratio_clamped_min"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_min"]) < 1.21' \ + 'min(data["train/probs_ratio_clamped_max"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_max"]) < 1.21' diff --git a/tests/functional/grpo_dp_simple.sh b/tests/functional/grpo_dp_simple.sh new file mode 100755 index 00000000000..a9611ad0264 --- /dev/null +++ b/tests/functional/grpo_dp_simple.sh @@ -0,0 +1,52 @@ +#!/bin/bash +# Lightweight e2e for grpo_sync.py ā€” TQ pipeline with the in-Ray "simple" +# backend. Same shape as tests/functional/grpo.sh (Qwen3-0.6B, 2 GPUs, +# 2 steps); flipping data_plane.enabled=true routes examples/run_grpo.py +# to nemo_rl.algorithms.grpo_sync.grpo_train_sync. + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) +# Mark the current repo as safe, since wandb fetches metadata about the repo +git config --global --add safe.directory $PROJECT_ROOT + +set -eou pipefail + +EXP_NAME=$(basename $0 .sh) +EXP_DIR=$SCRIPT_DIR/$EXP_NAME +LOG_DIR=$EXP_DIR/logs +JSON_METRICS=$EXP_DIR/metrics.json +RUN_LOG=$EXP_DIR/run.log +export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} + +rm -rf $EXP_DIR $LOG_DIR +mkdir -p $EXP_DIR $LOG_DIR + +cd $PROJECT_ROOT +uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ + $PROJECT_ROOT/examples/run_grpo.py \ + policy.model_name=Qwen/Qwen3-0.6B \ + grpo.num_prompts_per_step=2 \ + grpo.num_generations_per_prompt=4 \ + policy.train_global_batch_size=4 \ + policy.train_micro_batch_size=1 \ + cluster.gpus_per_node=2 \ + grpo.max_num_steps=2 \ + logger.tensorboard_enabled=true \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=false \ + logger.monitor_gpus=true \ + checkpointing.enabled=false \ + data_plane.enabled=true \ + data_plane.impl=transfer_queue \ + data_plane.backend=simple \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +uv run tests/check_metrics.py $JSON_METRICS \ + 'max(data["train/gen_kl_error"]) < 0.002' \ + 'min(data["train/probs_ratio_clamped_min"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_min"]) < 1.21' \ + 'min(data["train/probs_ratio_clamped_max"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_max"]) < 1.21' diff --git a/tests/unit/data_plane/README.md b/tests/unit/data_plane/README.md new file mode 100644 index 00000000000..df02430216d --- /dev/null +++ b/tests/unit/data_plane/README.md @@ -0,0 +1,222 @@ +# `tests/unit/data_plane/` ā€” test inventory + +Generated audit of every test function under `tests/unit/data_plane/` with a one-line summary. Use this when deciding what to consolidate or drop. + +--- + +## `test_architecture_invariants.py` (11 tests) + +- `test_grpo_sync_engages_tq_policy` ā€” Sync trainer must require a TQ-mediated policy. +- `test_grpo_sync_requires_data_plane_enabled` ā€” Sync trainer hard-fails when invoked without `data_plane.enabled=true`. +- `test_no_feature_gate_pattern_in_either_trainer` ā€” Catch the next "just one if branch" temptation in either trainer. +- `test_factory_does_not_construct_noop` ā€” Production factory must not return a NoOp client. +- `test_factory_rejects_disabled_impl` ā€” Factory must raise ā€” not return None / NoOp ā€” when disabled. +- `test_run_grpo_dispatches_both_trainers` ā€” `examples/run_grpo.py._select_trainer` returns legacy vs sync per config. +- `test_legacy_does_not_import_sync` ā€” Dependency direction: `grpo_sync.py` imports from `grpo.py`, not the reverse. +- `test_pack_per_token_field_is_exported` ā€” `pack_per_token_field` must be importable from `nemo_rl.data_plane.codec`. +- `test_pack_per_token_field_is_wired_into_writeback` ā€” **xfail.** At least one write-back call site must import it (wiring incomplete). +- `test_abc_method_present` ā€” Renaming an ABC method is a wire break ā€” keep the swap surface stable. +- `test_fp8_calib_filter_then_seqlen_check_no_crash` ā€” End-to-end behavioral repro of the job 11920261 calib-vs-seqlen bug. + +## `test_codec_jagged.py` (9 tests) + +- `test_to_nested_by_length_strips_padding` ā€” Right-pad columns must NOT be in the nested output. +- `test_to_nested_by_length_preserves_dtype` ā€” bf16 in ā†’ bf16 out. +- `test_to_nested_by_length_rejects_shape_mismatch` ā€” Shape sanity guard. +- `test_to_nested_by_length_rejects_1d_input` ā€” 1D inputs aren't valid (no seq dim). +- `test_materialize_pads_nested_with_field_specific_pad_value` ā€” Token field padded with pad_token_id; mask padded with 0. +- `test_materialize_passes_through_rectangular_tensors` ā€” Already-padded fields emitted unchanged. +- `test_materialize_jagged_layout_passes_nested_through` ā€” `layout='jagged'` path for nested-consuming callers. +- `test_materialize_default_pad_value_is_zero` ā€” No `pad_value_dict` ā†’ pad with 0. +- `test_response_from_nested_extracts_response_slice` ā€” Worker write-back: jagged (prompt+response) ā†’ response only. + +## `test_codec_mooncake.py` (4 tests) + +- `test_promote_1d_leaves_unsqueezes_1d` ā€” `_promote_1d_leaves` turns 1D `(N,)` leaves into `(N, 1)` for mooncake wire. +- `test_promote_1d_roundtrip_via_from_wire` ā€” `_promote_1d_leaves` + `_from_wire` restores original `(N,)` shape and values. +- `test_pack_per_token_field_truncates_sp_padding` ā€” pack_per_token_field slices each row to its own length, dropping SP padding. +- `test_pack_per_token_field_exact_fit_equals_maybe_pack_jagged` ā€” At exact fit, `pack_per_token_field` ā‰” `maybe_pack_jagged`. + +## `test_codec_wire_stripped.py` (5 tests) + +- `test_unwrap_wire_stripped_payload_empty_td_to_none` ā€” Empty TD (batch_dims=0) ā†’ None. +- `test_unwrap_wire_stripped_payload_real_nontensor_data_passes_through` ā€” Live NonTensorData payload survives unwrap. +- `test_materialize_handles_wire_stripped_nontensor_stack` ā€” Stack of empty TDs materializes to object array of None. +- `test_materialize_preserves_real_nontensor_data` ā€” NonTensorStack of strings materializes to raw strings. +- `test_materialize_decodes_nontensor_stack_with_tensor_field` ā€” Per-field decode: tensors stay padded, objects ride. + +## `test_correctness.py` (16 tests) + +- `test_kv_batch_get_after_clear_raises` ā€” v3 driver tried to read input_ids for log_data after clear ā€” must fail loud. +- `test_kv_batch_get_unproduced_field_raises` ā€” Requesting an unproduced field must raise, not return junk. +- `test_get_data_without_select_fields_raises` ā€” P2 invariant ā€” never silently fetch all fields. +- `test_kv_batch_put_rejects_non_tensor_leaves` ā€” P3 ā€” adapters reject non-tensor leaves; no pickle on the bus. +- `test_claim_meta_unregistered_task_raises` ā€” Catches typo'd consumer task names early. +- `test_kv_clear_with_none_drops_partition` ā€” Step-end teardown removes the partition entirely. +- `test_double_register_partition_is_idempotent_overwrite` ā€” Re-registering same partition_id within a step is OK. +- `test_check_consumption_status_only_true_when_all_consumed` ā€” Stage-done signal must not lie. +- `test_shard_meta_for_dp_partitions_keys_disjointly` ā€” Sum of shard sizes == total; pairwise disjoint. +- `test_shard_meta_for_dp_keeps_partition_id` ā€” partition_id propagated to every shard. +- `test_kv_first_write_carries_multimodal_extras_through_tq` ā€” VLM image features round-trip via TQ end-to-end. +- `test_kv_batch_put_preserves_bf16_dtype` ā€” Catches silent fp32 promotion. +- `test_kv_batch_put_preserves_int64_dtype` ā€” input_ids stays int64. +- `test_write_columns_accepts_batched_data_dict_input` ā€” Job 11614968 v2 crash guard: worker write-back accepts BatchedDataDict. +- `test_kv_first_write_rejects_key_count_mismatch` ā€” `len(keys) != n_samples` must fail (silent mis-align otherwise). +- `test_kv_first_write_meta_sequence_lengths_match_input_lengths` ā€” Megatron's balanced packing needs `meta.sequence_lengths` to match. + +## `test_factory.py` (5 tests) + +- `test_factory_none_cfg_rejected` ā€” None config fails fast, not silently. +- `test_factory_disabled_rejected` ā€” Production factory rejects disabled config. +- `test_factory_noop_impl_rejected` ā€” NoOp impl not selectable from production factory. +- `test_factory_unknown_impl_rejected` ā€” Unknown impl name fails fast with a helpful error. +- `test_factory_disabled_error_message_helpful` ā€” Disabled-config error message names the missing flag. + +## `test_interface_contract.py` (7 tests) + +- `test_factory_disabled_raises` ā€” Factory has no NoOp fallback ā€” disabled must raise. +- `test_factory_unknown_impl_raises` ā€” Unknown impl raises. +- `test_register_put_get_clear` ā€” End-to-end ABC round-trip. +- `test_claim_meta_advances_consumption` ā€” `claim_meta` advances the per-task consumption cursor. +- `test_get_data_requires_field_selection` ā€” P2 ā€” fetching all fields is forbidden. +- `test_kv_batch_put_rejects_non_tensor_leaves` ā€” P3 ā€” adapter rejects non-tensor leaves. +- `test_close_is_idempotent` ā€” `close()` can be called twice safely. + +## `test_kvbatchmeta.py` (10 tests) + +- `test_size_matches_keys` ā€” `size` derived from `sample_ids` length. +- `test_default_fields_and_extra_info_optional` ā€” `fields` and `sequence_lengths` default to None. +- `test_pickle_roundtrip_structural_equality` ā€” Cloudpickle round-trip for Ray actor dispatch. +- `test_keys_with_duplicates_allowed_or_warned` ā€” Meta doesn't enforce key uniqueness (caller's contract). +- `test_empty_meta_is_valid` ā€” Empty meta is a valid value (e.g. empty DP shard). +- `test_partition_id_is_required` ā€” `partition_id` is positional + required. +- `test_extra_info_default_is_unique_per_instance` ā€” Mutable default trap ā€” two metas don't share `extra_info`. +- `test_tags_align_with_keys` ā€” `tags` exactly one dict per key, or None. +- `test_tags_travel_with_subset_slice_concat` ā€” Per-key tags follow keys through subset/slice/concat. +- `test_tags_none_when_either_side_missing_in_concat` ā€” concat drops tags if either side has none. + +## `test_leader_broadcast.py` (2 tests) + +- `test_leader_broadcast_round_trip` ā€” 2-rank gloo broadcast of a BatchedDataDict round-trips. +- `test_get_replica_group_default_is_none` ā€” `TQWorkerMixin._get_replica_group` default is None. + +## `test_local_node_ip.py` (5 tests) + +- `test_local_node_ip_skips_link_local` ā€” gethostbyname returns 169.254.x.x ā†’ helper falls back. +- `test_local_node_ip_skips_loopback` ā€” Returns 127.0.0.1 ā†’ helper falls back. +- `test_local_node_ip_returns_routable` ā€” Routable address returned as-is. +- `test_local_node_ip_returns_empty_on_exception` ā€” DNS exception ā†’ returns empty string (no crash). +- `test_mc_tcp_bind_address_overwrites_existing` ā€” TQDataPlaneClient `__init__` uses direct assignment (not `setdefault`). + +## `test_message_log_decompose.py` (11 tests) + +- `test_decompose_message_log_basic_shapes` ā€” Basic shapes of decompose output. +- `test_decompose_message_log_no_assistant_turn` ā€” No-assistant case handled. +- `test_decompose_message_log_picks_first_assistant` ā€” Multiple assistant turns ā†’ first wins for `response_token_lengths`. +- `test_decompose_message_log_jagged_turn_count` ā€” Different turn counts pad `turn_lengths` with zeros. +- `test_decompose_message_log_missing_role_raises` ā€” Missing `role` raises KeyError loudly. +- `test_reconstruct_message_log_roundtrip` ā€” decompose ā†’ flatten ā†’ reconstruct equivalent message_log. +- `test_reconstruct_message_log_returns_views` ā€” Per-turn `token_ids` are views into local storage. +- `test_reconstruct_message_log_attaches_generation_logprobs` ā€” Attached only to assistant turns. +- `test_attach_message_log_view_populates_batch` ā€” `attach_message_log_view` populates batch view. +- `test_attach_message_log_view_noop_when_fields_absent` ā€” Without decomposed fields, attach is a no-op. +- `test_attach_message_log_view_idempotent` ā€” Calling twice produces same shape. + +## `test_observability.py` (8 tests) + +- `test_put_records_bytes_and_count` ā€” Observability decorator records put bytes + count. +- `test_get_records_after_put` ā€” Records get ops after put. +- `test_register_and_clear_recorded` ā€” register/clear ops are recorded. +- `test_error_status_recorded_and_reraised` ā€” Decorator records error AND re-raises (no swallowing). +- `test_snapshot_accumulates_successful_ops` ā€” Snapshot accumulates over time. +- `test_default_callback_is_noop` ā€” Omitting on_event must not raise. +- `test_close_propagates` ā€” close() is forwarded to wrapped client. +- `test_factory_wraps_when_observability_enabled` ā€” factory.py uses the same MetricsDataPlaneClient. + +## `test_preshard_extras.py` (10 tests) + +- `test_kv_first_write_writes_seed_fields` ā€” Seed fields written to TQ. +- `test_kv_first_write_carries_multimodal_extras` ā€” VLM extras (pixel_values) ride along, no schema declaration needed. +- `test_kv_first_write_keys_match_uids_x_ngen` ā€” Keys round-trip: `f"{uid}_g{i}"` preserved. +- `test_shard_meta_for_dp_partitions_keys_disjointly` ā€” Sum of shards == total, disjoint. +- `test_shard_meta_for_dp_preserves_partition_id` ā€” partition_id preserved across DP shards. +- `test_shard_meta_for_dp_unsorted_round_trip` ā€” `unsorted_indices` reconstructs input order from concat. +- `test_kvbatchmeta_subset_filters_keys_and_seqlens` ā€” `subset` filters keys + seq_lengths. +- `test_kvbatchmeta_concat_joins_keys_and_seqlens` ā€” `concat` joins. +- `test_kvbatchmeta_slice_takes_range` ā€” `slice` takes a contiguous range. +- `test_kvbatchmeta_concat_rejects_partition_mismatch` ā€” `concat` rejects different `partition_id`s. + +## `test_seqpack_equivalence.py` (3 tests, Ɨ2 backends) + +- `test_seqpack_legacy_equals_tq[simple|mooncake_cpu]` ā€” Sequence packing byte-equivalence: legacy shards == TQ-roundtripped. +- `test_dynbatch_legacy_equals_tq[simple|mooncake_cpu]` ā€” Same claim for dynamic batching. +- `test_no_packing_legacy_equals_tq[simple|mooncake_cpu]` ā€” Sanity: lossless transport even without packing/dynbatch. + +## `test_smoke.py` (5 tests) + +- `test_sync_utils_module_imports` ā€” Catches FQN drift after `algorithms.sync_utils` consolidation. +- `test_sync_rollout_actor_registered_under_vllm_tier` ā€” Multinode dep: tensordict must be on the vLLM tier. +- `test_kvbatchmeta_schema_unchanged` ā€” Schema-pin: KVBatchMeta is the cross-process boundary. +- `test_dataplane_client_abc_surface` ā€” Catches accidental ABC method removal/rename. +- `test_async_and_sync_actors_share_env_tier` ā€” Sync mirrors async's env tier (both drive vLLM). + +## `test_sync_one_hop.py` (9 tests) + +- `test_write_columns_lands_in_tq` ā€” write_columns lands fields in TQ. +- `test_read_columns_returns_only_requested_fields` ā€” read_columns honors `select_fields`. +- `test_write_then_read_roundtrip_after_train_window` ā€” Full lifecycle: rollout puts ā†’ driver deltas ā†’ read deltas back. +- `test_meta_keys_identity_across_dp_shards` ā€” `shard_meta_for_dp` must NOT mint new keys. +- `test_kv_clear_uses_meta_keys_minted_at_rollout` ā€” Step-end clear targets the SAME keys rollout minted. +- `test_apply_dynamic_sampling_filters_zero_std` ā€” Drops zero-std uids and clears their TQ payload. +- `test_apply_dynamic_sampling_completes_when_train_size_reached` ā€” When cache hits train_prompts_size, is_complete=True. +- `test_apply_dynamic_sampling_overflow_slices_and_clears` ā€” Overflow: slice + clear discards. +- `test_apply_dynamic_sampling_raises_on_max_gen_batches` ā€” Exceeding max_gen_batches raises loudly. + +## `test_tq_lifecycle.py` (5 tests, some Ɨ2 backends) + +- `test_smoke_round_trip` ā€” Basic register ā†’ put ā†’ claim_meta ā†’ get_data ā†’ clear flow. +- `test_smoke_round_trip_backends[simple|mooncake_cpu]` ā€” Same parameterized over both backends. +- `test_smoke_round_trip_1d_fields` ā€” `(N,)` tensors come back as `(N,)`, not `(N,1)`. +- `test_object_round_trip_backends[simple|mooncake_cpu]` ā€” `np.ndarray(dtype=object)` round-trips both backends. +- `test_object_and_tensor_mixed_round_trip_backends[simple|mooncake_cpu]` ā€” Mixed tensor+object in one put. + +--- + +## Potential simplifications (candidates to drop or merge) + +| Overlap | Files involved | Suggestion | +|---|---|---| +| `factory disabled/unknown impl rejected` | `test_factory.py` (5 tests) + `test_interface_contract.py::test_factory_*` (2) | Keep `test_factory.py` (more thorough); drop the two duplicates in `test_interface_contract.py` | +| `kv_batch_put_rejects_non_tensor_leaves` | `test_correctness.py` + `test_interface_contract.py` | One is enough ā€” keep `test_correctness.py`'s (P3 framing). | +| `get_data_without_select_fields_raises` / `test_get_data_requires_field_selection` | `test_correctness.py` + `test_interface_contract.py` | Same property; keep `test_correctness.py`. | +| `shard_meta_for_dp_partitions_keys_disjointly` + `_keeps/preserves_partition_id` | `test_correctness.py` (2) + `test_preshard_extras.py` (2) | Pure dup. Drop from `test_correctness.py`. | +| `kv_first_write_carries_multimodal_extras` | `test_correctness.py::test_kv_first_write_carries_multimodal_extras_through_tq` + `test_preshard_extras.py::test_kv_first_write_carries_multimodal_extras` | Pure dup. Keep `test_preshard_extras.py`. | +| ABC surface checks | `test_smoke.py::test_dataplane_client_abc_surface` + `test_architecture_invariants.py::test_abc_method_present` + `test_interface_contract.py` (covers same surface end-to-end) | Three angles on the same invariant. Keep `test_architecture_invariants.py` (most explicit); drop the smoke one. | +| Codec tests across 3 files | `test_codec_jagged.py` (9), `test_codec_mooncake.py` (4), `test_codec_wire_stripped.py` (5) | Distinct paths but small files ā€” could merge into a single `test_codec.py` with `# ā”€ā”€ jagged ā”€ā”€` / `# ā”€ā”€ mooncake ā”€ā”€` / `# ā”€ā”€ wire_stripped ā”€ā”€` sections. Saves 2 file headers. | +| `test_smoke.py` ā€” 5 narrow checks | various | These are best as a single fast "import-this-stuff" smoke test, not 5 separate ones. Consider folding into a parametrized `test_imports_unchanged`. | + +### Likely to drop + +If you want a one-pass cull, the safest deletes are: +1. `test_interface_contract.py::test_factory_disabled_raises` (dup of `test_factory.py`) +2. `test_interface_contract.py::test_factory_unknown_impl_raises` (dup of `test_factory.py`) +3. `test_interface_contract.py::test_get_data_requires_field_selection` (dup of `test_correctness.py`) +4. `test_interface_contract.py::test_kv_batch_put_rejects_non_tensor_leaves` (dup of `test_correctness.py`) +5. `test_correctness.py::test_shard_meta_for_dp_partitions_keys_disjointly` (dup of `test_preshard_extras.py`) +6. `test_correctness.py::test_shard_meta_for_dp_keeps_partition_id` (dup of `test_preshard_extras.py`) +7. `test_correctness.py::test_kv_first_write_carries_multimodal_extras_through_tq` (dup of `test_preshard_extras.py`) +8. `test_smoke.py::test_dataplane_client_abc_surface` (dup of `test_architecture_invariants.py`) + +ā†’ āˆ’8 tests, no coverage loss. + +### Likely to consolidate (file count, not test count) + +- Merge `test_codec_{jagged,mooncake,wire_stripped}.py` ā†’ `test_codec.py` (3 files ā†’ 1, same 18 tests) +- Merge `test_factory.py` into `test_interface_contract.py` (or vice-versa) since they share scope +- `test_smoke.py` is just 5 import/registration checks ā€” could move into `test_architecture_invariants.py` + +### File-count target + +| Now | After dedupe + merge | +|---|---| +| 17 files / ~125 tests | 12 files / ~117 tests | diff --git a/tests/unit/data_plane/__init__.py b/tests/unit/data_plane/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/data_plane/_rollout_shapes.py b/tests/unit/data_plane/_rollout_shapes.py new file mode 100644 index 00000000000..3bb2e614522 --- /dev/null +++ b/tests/unit/data_plane/_rollout_shapes.py @@ -0,0 +1,244 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Realistic rollout-shaped data builders + shared test helpers. + +Mints data with the *shape and types* an actual GRPO rollout produces ā€” +mixed dtypes (bf16 logprobs, int64 ids, int32 masks), realistic value +distributions, optional multimodal extras, and varied multi-turn message +logs. Use these instead of inline toy tensors so tests cover the same +type / scenario complexity as production runs. + +Also exposes a handful of small cross-file test helpers (uid ā†’ key +minting, TQ partition setup, mooncake availability) that several test +files used to duplicate. + +Helpers are plain functions (not pytest fixtures) so they're explicit at +the call site and don't depend on a conftest. +""" + +from __future__ import annotations + +import os +from typing import Any + +import numpy as np +import torch + +from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient +from nemo_rl.data_plane.schema import DP_TRAIN_FIELDS + + +def make_rollout_batch( + n: int = 8, + max_seqlen: int = 256, + *, + multimodal: bool = False, + logprob_dtype: torch.dtype = torch.bfloat16, + id_dtype: torch.dtype = torch.long, + mask_dtype: torch.dtype = torch.int32, + seed: int = 42, +) -> dict[str, Any]: + """Return a fields dict shaped like rollout's first put. + + Mirrors what ``SyncRolloutActor.rollout_to_tq`` actually writes: + int64 token ids, int32 masks, bf16 logprobs, fp32 (or bf16) advantages, + optional ``multi_modal_inputs`` dict for VLM models. + + Args: + n: Batch size. + max_seqlen: Padded sequence length; per-row valid length is in + ``[max_seqlen//4, max_seqlen]``. + multimodal: Include ``multi_modal_inputs`` (pixel_values + image_grid_thw). + logprob_dtype: dtype for logprobs/advantages (real runs use bf16). + id_dtype: dtype for input_ids/lengths. + mask_dtype: dtype for masks. + seed: RNG seed for reproducibility. + + Returns: + Dict with keys ``input_ids``, ``input_lengths``, ``attention_mask``, + ``token_mask``, ``sample_mask``, ``generation_logprobs``, + ``prev_logprobs``, ``reference_policy_logprobs``, ``advantages``, + optionally ``multi_modal_inputs``. + """ + g = torch.Generator().manual_seed(seed) + + # Per-row valid lengths spanning ~25-100% of max_seqlen. + low = max(1, max_seqlen // 4) + lengths = torch.randint(low, max_seqlen + 1, (n,), generator=g).to(id_dtype) + + # Token ids: random vocab-shaped, padded with 0. + input_ids = torch.zeros((n, max_seqlen), dtype=id_dtype) + for i in range(n): + nrow = int(lengths[i]) + input_ids[i, :nrow] = torch.randint(1, 50000, (nrow,), generator=g) + + # Masks: 1 for valid tokens, 0 for padding. + token_mask = torch.zeros((n, max_seqlen), dtype=mask_dtype) + for i in range(n): + token_mask[i, : int(lengths[i])] = 1 + attention_mask = token_mask.clone() + sample_mask = torch.ones((n,), dtype=mask_dtype) + + # Logprobs: realistic distribution centered around -2.0 (typical token logprob), + # std ~1 ā€” catches dtype-narrowing bugs that pass on zero inputs. + def _lp() -> torch.Tensor: + return (torch.randn(n, max_seqlen, generator=g) - 2.0).to(logprob_dtype) + + out: dict[str, Any] = { + "input_ids": input_ids, + "input_lengths": lengths.to(torch.long), + "attention_mask": attention_mask, + "token_mask": token_mask, + "sample_mask": sample_mask, + "generation_logprobs": _lp(), + "prev_logprobs": _lp(), + "reference_policy_logprobs": _lp(), + "advantages": torch.randn(n, max_seqlen, generator=g).to(logprob_dtype), + } + + if multimodal: + # VLM extras as flat top-level fields (the codec wire format ā€” + # nested dicts aren't valid leaves). Real production writes these + # with similar shapes; we keep them small for fast tests. + T, H, W = 1, 8, 8 + n_image_tokens = T * H * W + out["pixel_values"] = torch.randn(n, n_image_tokens, 3, generator=g).to( + torch.bfloat16 + ) + out["image_grid_thw"] = torch.tensor([[T, H, W]] * n, dtype=torch.long) + + return out + + +def make_realistic_tags( + n: int, + *, + zero_std_fraction: float = 0.25, + seed: int = 42, +) -> list[dict[str, float | int]]: + """Per-sample tags as produced by the GRPO driver after baseline/std compute. + + Mirrors what gets stamped onto ``KVBatchMeta.tags`` for dynamic-sampling + filtering. Some rows have ``std=0.0`` (zero-variance, filtered) and + others non-zero (survivors). + + Args: + n: Number of samples. + zero_std_fraction: Fraction of rows tagged with ``std=0.0`` (filtered). + seed: RNG seed. + """ + rng = np.random.default_rng(seed) + n_zero = int(round(n * zero_std_fraction)) + stds = np.concatenate([np.zeros(n_zero), rng.uniform(0.1, 1.5, size=n - n_zero)]) + rng.shuffle(stds) + rewards = rng.uniform(-1.0, 1.0, size=n) + prompt_ids = rng.integers(0, 1000, size=n) + return [ + { + "std": float(stds[i]), + "total_reward": float(rewards[i]), + "prompt_id": int(prompt_ids[i]), + "weight_version": 1, + } + for i in range(n) + ] + + +def make_multi_turn_message_log( + n: int, + *, + turns_per_sample: list[int] | None = None, + seed: int = 42, +) -> list[list[dict[str, Any]]]: + """Realistic multi-turn message_log: list-of-turn-dicts per sample. + + Each turn carries ``role`` (alternating user/assistant), ``content`` + (string), and ``token_ids`` (int64 tensor). Variable turn counts + capture the jagged case that ``decompose_message_log`` flattens. + + Args: + n: Number of samples. + turns_per_sample: Optional explicit turn count per sample. If + None, random in ``[1, 4]``. + seed: RNG seed. + """ + g = torch.Generator().manual_seed(seed) + if turns_per_sample is None: + turns_per_sample = [int(t) for t in torch.randint(1, 5, (n,), generator=g)] + out: list[list[dict[str, Any]]] = [] + for i, k in enumerate(turns_per_sample): + sample_log: list[dict[str, Any]] = [] + for t in range(k): + role = "user" if t % 2 == 0 else "assistant" + tok_len = int(torch.randint(8, 64, (1,), generator=g)) + sample_log.append( + { + "role": role, + "content": f"sample_{i}_turn_{t}_text", + "token_ids": torch.randint( + 1, 50000, (tok_len,), generator=g, dtype=torch.long + ), + } + ) + out.append(sample_log) + return out + + +# ā”€ā”€ Cross-file test helpers (deduped from per-file definitions) ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def keys_from_uids(uids: list[str], n_gen: int = 1) -> list[str]: + """Mint per-generation sample keys from prompt uids: ``f"{uid}_g{i}"``. + + Mirrors the production rollout convention ā€” one key per generation per + prompt ā€” so tests share the same uid ā†’ key mapping as the trainer. + """ + return [f"{uid}_g{i}" for uid in uids for i in range(n_gen)] + + +def register_train_partition( + client: NoOpDataPlaneClient, + *, + num_samples: int, + fields: list[str] | None = None, + partition_id: str = "train", + consumer_tasks: list[str] | None = None, +) -> None: + """Open a TQ partition with the train-side defaults (``DP_TRAIN_FIELDS`` + ``["train"]``). + + Centralizes the boilerplate three test files used to inline as + ``_setup`` / ``_setup_partition``. + """ + client.register_partition( + partition_id=partition_id, + fields=list(fields if fields is not None else DP_TRAIN_FIELDS), + num_samples=num_samples, + consumer_tasks=consumer_tasks if consumer_tasks is not None else ["train"], + ) + + +def mooncake_available() -> bool: + """Return True if the ``mooncake`` wheel is importable. + + Set ``NEMO_RL_REQUIRE_MOONCAKE=1`` to promote a missing import into a + loud ``ImportError`` instead of returning False ā€” so CI fails when + the wheel is expected but absent. + """ + try: + import mooncake # noqa: F401 + except ImportError: + if os.environ.get("NEMO_RL_REQUIRE_MOONCAKE") == "1": + raise + return False + return True diff --git a/tests/unit/data_plane/test_architecture_invariants.py b/tests/unit/data_plane/test_architecture_invariants.py new file mode 100644 index 00000000000..dbd8fd0bac9 --- /dev/null +++ b/tests/unit/data_plane/test_architecture_invariants.py @@ -0,0 +1,72 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Minimal behavioral invariants for the data-plane wiring. + +* ``examples/run_grpo._select_trainer`` dispatches the legacy trainer + when ``data_plane`` is absent and the sync trainer when enabled. +* The ``DataPlaneClient`` ABC carries every method adapters depend on. +""" + +from __future__ import annotations + +import pathlib + +import pytest + +REPO = pathlib.Path(__file__).resolve().parents[3] + + +def test_run_grpo_dispatches_both_trainers(): + """``examples/run_grpo._select_trainer`` returns the TQ-mediated + ``grpo_train_sync`` iff ``data_plane.enabled`` is true, and the + legacy ``grpo_train`` otherwise.""" + import sys + + sys.path.insert(0, str(REPO / "examples")) + try: + from run_grpo import _select_trainer + finally: + sys.path.pop(0) + from nemo_rl.algorithms.grpo import MasterConfig, grpo_train + from nemo_rl.algorithms.grpo_sync import grpo_train_sync + + cfg_legacy = MasterConfig.model_construct(data_plane=None) + assert _select_trainer(cfg_legacy) is grpo_train + + cfg_sync = MasterConfig.model_construct(data_plane={"enabled": True}) + assert _select_trainer(cfg_sync) is grpo_train_sync + + +@pytest.mark.parametrize( + "method", + [ + "register_partition", + "claim_meta", + "get_data", + "put_samples", + "get_samples", + "clear_samples", + "check_consumption_status", + "close", + ], +) +def test_data_plane_client_abc_method_present(method: str) -> None: + """The ``DataPlaneClient`` ABC is the swap surface; a silent rename + is a breaking change for every adapter.""" + from nemo_rl.data_plane.interfaces import DataPlaneClient + + assert hasattr(DataPlaneClient, method), ( + f"DataPlaneClient ABC is missing required method {method!r}. " + "This is a breaking change for every adapter." + ) diff --git a/tests/unit/data_plane/test_codec_jagged.py b/tests/unit/data_plane/test_codec_jagged.py new file mode 100644 index 00000000000..9cbee1c38f3 --- /dev/null +++ b/tests/unit/data_plane/test_codec_jagged.py @@ -0,0 +1,233 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the padded ā†” jagged codec bridge. + +Phase 1 of the wire-jagged plan: writer emits nested, reader pads on +demand. These tests cover the conversion helpers in isolation; e2e +parity is validated separately. +""" + +from __future__ import annotations + +import pytest +import torch +from tensordict import TensorDict + +from nemo_rl.data_plane.codec import ( + materialize, + response_from_nested, + to_nested_by_length, +) + +from ._rollout_shapes import make_rollout_batch + + +def _padded(rows: list[list[int]], pad: int = 0) -> tuple[torch.Tensor, torch.Tensor]: + """Pad a list of int sequences to a rectangle; return (padded, lengths).""" + n = len(rows) + s = max(len(r) for r in rows) + out = torch.full((n, s), pad, dtype=torch.long) + lens = torch.tensor([len(r) for r in rows], dtype=torch.long) + for i, r in enumerate(rows): + out[i, : len(r)] = torch.tensor(r, dtype=torch.long) + return out, lens + + +# ā”€ā”€ to_nested_by_length ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_to_nested_by_length_strips_padding() -> None: + """The right-pad columns must NOT be in the nested output.""" + padded, lens = _padded([[1, 2, 3], [4, 5], [6, 7, 8, 9]], pad=0) + nested = to_nested_by_length(padded, lens) + assert nested.is_nested + rows = list(nested.unbind()) + assert torch.equal(rows[0], torch.tensor([1, 2, 3])) + assert torch.equal(rows[1], torch.tensor([4, 5])) + assert torch.equal(rows[2], torch.tensor([6, 7, 8, 9])) + + +def test_to_nested_by_length_preserves_dtype() -> None: + """bf16 in ā†’ bf16 out.""" + padded = torch.randn((3, 5), dtype=torch.bfloat16) + lens = torch.tensor([2, 4, 5], dtype=torch.long) + nested = to_nested_by_length(padded, lens) + assert nested.dtype == torch.bfloat16 + + +def test_to_nested_by_length_rejects_shape_mismatch() -> None: + padded = torch.zeros((3, 4)) + bad_lens = torch.tensor([1, 2]) # only 2, not 3 + with pytest.raises(ValueError, match=r"lengths shape"): + to_nested_by_length(padded, bad_lens) + + +def test_to_nested_by_length_rejects_1d_input() -> None: + with pytest.raises(ValueError, match=r"\(N, S"): + to_nested_by_length(torch.zeros(5), torch.tensor([5])) + + +# ā”€ā”€ materialize: jagged ā†’ padded ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_materialize_pads_nested_with_field_specific_pad_value() -> None: + """Token field padded with pad_token_id; mask padded with 0. + + This is the contract worker code expects: the padded view it + receives looks identical to a rectangular tensor produced by + batched_message_log_to_flat_message. + """ + ids_padded, lens = _padded([[10, 20, 30], [40, 50], [60, 70, 80, 90]], pad=0) + mask_padded, _ = _padded([[1, 1, 1], [1, 1], [1, 1, 1, 1]], pad=0) + ids_nested = to_nested_by_length(ids_padded, lens) + mask_nested = to_nested_by_length(mask_padded, lens) + + td = TensorDict( + {"input_ids": ids_nested, "token_mask": mask_nested}, + batch_size=[3], + ) + + bdd = materialize( + td, + layout="padded", + pad_value_dict={"input_ids": 999, "token_mask": 0}, + ) + + # Tokens are padded with the requested ID, not 0. + assert bdd["input_ids"].shape == (3, 4) + assert bdd["input_ids"][0, 3].item() == 999 # row 0 needs 1 pad + assert bdd["input_ids"][1, 2].item() == 999 # row 1 needs 2 pads + assert bdd["input_ids"][1, 3].item() == 999 + assert bdd["input_ids"][2, 3].item() == 90 # row 2 needs no padding + + # Mask uses the default 0 ā€” match the source. + assert bdd["token_mask"].shape == (3, 4) + assert bdd["token_mask"][0, 3].item() == 0 + assert bdd["token_mask"][2, 3].item() == 1 + + +def test_materialize_passes_through_rectangular_tensors() -> None: + """Already-padded fields are emitted unchanged (no spurious copy).""" + rect = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.long) + td = TensorDict({"sample_mask": rect}, batch_size=[2]) + bdd = materialize(td, layout="padded") + assert torch.equal(bdd["sample_mask"], rect) + + +def test_materialize_jagged_layout_passes_nested_through() -> None: + """``layout='jagged'`` is the path for callers that consume nested.""" + padded, lens = _padded([[1, 2], [3, 4, 5]], pad=0) + nested = to_nested_by_length(padded, lens) + td = TensorDict({"x": nested}, batch_size=[2]) + bdd = materialize(td, layout="jagged") + assert bdd["x"].is_nested + + +def test_materialize_default_pad_value_is_zero() -> None: + """No pad_value_dict ā†’ fields pad with 0.""" + padded, lens = _padded([[1, 2, 3], [4]], pad=0) + nested = to_nested_by_length(padded, lens) + td = TensorDict({"x": nested}, batch_size=[2]) + bdd = materialize(td, layout="padded") + assert bdd["x"][1, 1].item() == 0 + assert bdd["x"][1, 2].item() == 0 + + +# ā”€ā”€ response_from_nested ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_response_from_nested_extracts_response_slice() -> None: + """Worker write-back path: jagged (prompt+response) ā†’ response only. + + With the verl convention, output position i corresponds to predicting + input token i+1 ā€” so the slice is left-shifted by one. + """ + # Two samples: prompt_len=2, resp_len=3 / prompt_len=1, resp_len=2 + full_rows = [ + torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5]), # prompt 0,1; resp 2,3,4 + torch.tensor([1.1, 1.2, 1.3]), # prompt 0; resp 1,2 + ] + full = torch.nested.as_nested_tensor(full_rows, layout=torch.jagged) + resp_mask_rows = [ + torch.tensor([1.0, 1.0, 1.0]), # response_len = 3 + torch.tensor([1.0, 1.0]), # response_len = 2 + ] + response_mask = torch.nested.as_nested_tensor(resp_mask_rows, layout=torch.jagged) + + out = response_from_nested(full, response_mask) + assert out.is_nested + rows = list(out.unbind()) + # Row 0: full has 5 tokens; resp_len=3 ā†’ values[5-3-1:5-1] = values[1:4] = [0.2, 0.3, 0.4] + assert torch.allclose(rows[0], torch.tensor([0.2, 0.3, 0.4])) + # Row 1: full has 3 tokens; resp_len=2 ā†’ values[3-2-1:3-1] = values[0:2] = [1.1, 1.2] + assert torch.allclose(rows[1], torch.tensor([1.1, 1.2])) + + +# ā”€ā”€ Realistic-shape coverage using ``_rollout_shapes.make_rollout_batch`` ā”€ā”€ +# These exercise the same codec helpers with the exact dtypes + value +# distributions a real GRPO rollout produces (bf16 logprobs, int64 ids, +# int32 masks, variable per-row lengths). Catches dtype-narrowing and +# padding-arithmetic bugs that pass on the toy data above. + + +@pytest.mark.parametrize( + "logprob_dtype", + [torch.bfloat16, torch.float32], + ids=["bf16", "fp32"], +) +def test_to_nested_by_length_realistic_logprobs(logprob_dtype: torch.dtype) -> None: + """``generation_logprobs`` shape (bf16/fp32) from a real rollout shape round-trips.""" + + batch = make_rollout_batch(n=8, max_seqlen=128, logprob_dtype=logprob_dtype, seed=7) + nested = to_nested_by_length(batch["generation_logprobs"], batch["input_lengths"]) + # dtype must survive the conversion (bf16 in ā†’ bf16 out). + assert nested.dtype == logprob_dtype + # Per-row valid region matches the input. + for i, row in enumerate(nested.unbind()): + valid = int(batch["input_lengths"][i]) + assert row.shape[0] == valid + assert torch.equal( + row, batch["generation_logprobs"][i, :valid].to(logprob_dtype) + ) + + +def test_materialize_realistic_full_field_set_preserves_dtypes() -> None: + """All rollout fields round-trip through ``materialize`` with correct dtypes. + + Catches the class of bugs where padding silently upcasts bf16 ā†’ fp32 or + coerces int64 ā†’ int32 because pad_value_dict's defaults were the wrong type. + """ + + batch = make_rollout_batch(n=4, max_seqlen=64, seed=11) + # Build a wire TD with jagged leaves keyed by field name. + td = TensorDict( + { + "input_ids": to_nested_by_length( + batch["input_ids"], batch["input_lengths"] + ), + "generation_logprobs": to_nested_by_length( + batch["generation_logprobs"], batch["input_lengths"] + ), + "token_mask": to_nested_by_length( + batch["token_mask"], batch["input_lengths"] + ), + }, + batch_size=[4], + ) + out = materialize(td, layout="padded", pad_value_dict={"input_ids": 0}) + + # Each field comes back at its original dtype. + assert out["input_ids"].dtype == torch.long + assert out["generation_logprobs"].dtype == torch.bfloat16 + assert out["token_mask"].dtype == torch.int32 diff --git a/tests/unit/data_plane/test_codec_mooncake.py b/tests/unit/data_plane/test_codec_mooncake.py new file mode 100644 index 00000000000..469cd703c01 --- /dev/null +++ b/tests/unit/data_plane/test_codec_mooncake.py @@ -0,0 +1,154 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the mooncake_cpu-specific wire workarounds. + +Covers: + P1 ā€” `promote_1d` round-trip: writer unsqueezes 1D ā†’ (N,1), reader squeezes back. + P2 ā€” pack_per_token_field: tolerates SP padding wider than max(lengths). + +No Ray, no GPU, no transfer_queue required. +""" + +from __future__ import annotations + +import torch + +from nemo_rl.data_plane.codec import pack_per_token_field + +from ._rollout_shapes import make_rollout_batch + +# ā”€ā”€ P1: promote_1d ā€” writer unsqueezes, reader squeezes ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_promote_1d_leaves_unsqueezes_1d() -> None: + """`_promote_1d_leaves` turns 1D ``(N,)`` leaves into ``(N, 1)``. + + Guards the mooncake_cpu path where TQ's extract_field_schema silently + unsqueezes 1D fields in metadata; the wire layer pre-unsqueezes so the + per-row data shape matches the metadata-recorded shape. + """ + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.transfer_queue import _promote_1d_leaves + + n = 8 + t = torch.arange(n, dtype=torch.float32) + td = TensorDict({"reward": t}, batch_size=[n]) + + out = _promote_1d_leaves(td) + assert out["reward"].shape == (n, 1), ( + f"Expected wire shape ({n}, 1) but got {tuple(out['reward'].shape)}." + ) + + +def test_promote_1d_roundtrip_via_from_wire() -> None: + """`_promote_1d_leaves` then `_from_wire` restores the original ``(N,)`` shape and values.""" + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.transfer_queue import ( + _from_wire, + _promote_1d_leaves, + ) + + n = 6 + original = torch.arange(n, dtype=torch.float32) + td = TensorDict({"reward": original}, batch_size=[n]) + + wire = _promote_1d_leaves(td) + assert wire["reward"].shape == (n, 1) + + back = _from_wire(wire) + assert back["reward"].shape == (n,) + assert torch.equal(back["reward"], original) + + +# ā”€ā”€ P2: pack_per_token_field ā€” tolerates SP padding ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_pack_per_token_field_truncates_sp_padding() -> None: + """pack_per_token_field slices each row to its own length, dropping SP padding. + + mcore SP rounds the forward output's seq dim up to a multiple of TP, so + val.shape[1] > max(lengths). maybe_pack_jagged would skip this field + (wrong shape); pack_per_token_field handles it correctly. + """ + + n, max_len, sp_extra = 4, 8, 3 # val is wider by sp_extra tokens + lengths = torch.tensor([3, 5, 7, 4], dtype=torch.long) + assert lengths.max().item() == max_len - 1 # max_len=8 > max(lengths)=7 + val = torch.randn(n, max_len + sp_extra) # (4, 11) + + out = pack_per_token_field(val, lengths) + + assert out.is_nested, "pack_per_token_field must produce a nested tensor." + rows = list(out.unbind()) + assert len(rows) == n + for i, row in enumerate(rows): + expected_len = int(lengths[i].item()) + assert row.shape == (expected_len,), ( + f"Row {i}: expected length {expected_len}, got {tuple(row.shape)}. " + "SP padding tail was not dropped." + ) + assert torch.equal(row, val[i, :expected_len]), ( + f"Row {i}: values differ after truncation." + ) + + +def test_pack_per_token_field_exact_fit_equals_maybe_pack_jagged() -> None: + """When val.shape[1] == max(lengths), pack_per_token_field and + maybe_pack_jagged produce identical jagged outputs. + + This is the 'no SP padding' case ā€” the two helpers must agree when + the input is already exactly the right width. + """ + from nemo_rl.data_plane.codec import maybe_pack_jagged, pack_per_token_field + + n = 4 + lengths = torch.tensor([3, 5, 2, 4], dtype=torch.long) + max_len = int(lengths.max().item()) + val = torch.randn(n, max_len) + + out_pack = pack_per_token_field(val, lengths) + out_maybe = maybe_pack_jagged(val, lengths) + + assert out_pack.is_nested + assert out_maybe.is_nested + + rows_pack = list(out_pack.unbind()) + rows_maybe = list(out_maybe.unbind()) + for i, (rp, rm) in enumerate(zip(rows_pack, rows_maybe)): + assert torch.equal(rp, rm), ( + f"Row {i} differs between pack_per_token_field and maybe_pack_jagged " + "on an exact-fit input." + ) + + +# ā”€ā”€ Realistic bf16 per-token coverage ā”€ā”€ + + +def test_pack_per_token_field_realistic_bf16_logprobs() -> None: + """pack_per_token_field on bf16 prev_logprobs (realistic dtype + value distribution).""" + + batch = make_rollout_batch( + n=6, max_seqlen=96, logprob_dtype=torch.bfloat16, seed=29 + ) + out = pack_per_token_field(batch["prev_logprobs"], batch["input_lengths"]) + assert out.is_nested + assert out.dtype == torch.bfloat16 + # Per-row valid region matches input ā€” bf16 round-trip is loss-y at the bit + # level but pack_per_token_field shouldn't change values. + for i, row in enumerate(out.unbind()): + valid = int(batch["input_lengths"][i]) + assert row.shape[0] == valid + assert torch.equal(row, batch["prev_logprobs"][i, :valid]) diff --git a/tests/unit/data_plane/test_codec_wire_stripped.py b/tests/unit/data_plane/test_codec_wire_stripped.py new file mode 100644 index 00000000000..1646e33d528 --- /dev/null +++ b/tests/unit/data_plane/test_codec_wire_stripped.py @@ -0,0 +1,190 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Regression tests for the wire-stripped ``NonTensorStack`` path. + +TQ's simple-backend ``MsgpackEncoder._encode_tensordict`` serializes any +``TensorDictBase`` via ``dict(obj.items())`` ā€” only the tensor backing +dict. ``NonTensorData`` stores its payload in ``_non_tensordict["data"]``, +so it round-trips through ZMQ as an empty +``TensorDict({}, batch_size=[])`` ā€” the string payload is silently +dropped. The simple-backend storage manager's ``_pack_field_values`` +then assembles those stripped TDs into a ``NonTensorStack`` that +``materialize`` has to defend against. The pre-fix path crashed with +``RuntimeError: generator raised StopIteration``. + +Construction note: ``tensordict>=0.12.2`` rejects +``NonTensorStack(TensorDict({}, batch_size=[]), ...)`` at construction +time (``All tensordicts must be non-tensors``). To validate +``materialize``'s decode without skirting tensordict's invariants we: + +* test :func:`unwrap_wire_stripped_payload` directly ā€” pure per-item + helper, accepts the wire-stripped ``TensorDict`` shape without + needing the stack constructor at all; +* drive :func:`materialize` end-to-end by patching ``.tolist()`` on a + constructed (valid) ``NonTensorStack`` so it returns the wire-stripped + items list ā€” preserves the data-in / data-out contract while routing + around the constructor's homogeneity check. +""" + +from __future__ import annotations + +from unittest.mock import patch + +import numpy as np +import torch +from tensordict import NonTensorData, NonTensorStack, TensorDict + +from nemo_rl.data.llm_message_utils import decompose_message_log +from nemo_rl.data_plane.codec import ( + materialize, + to_nested_by_length, + unwrap_wire_stripped_payload, +) + +from ._rollout_shapes import make_multi_turn_message_log + +# ā”€ā”€ unwrap_wire_stripped_payload ā€” direct per-item coverage ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_unwrap_wire_stripped_payload_empty_td_to_none() -> None: + """An empty ``TensorDict`` (batch_dims=0, no keys) ā†’ ``None``.""" + assert unwrap_wire_stripped_payload(TensorDict({}, batch_size=[])) is None + + +def test_unwrap_wire_stripped_payload_real_nontensor_data_passes_through() -> None: + """A live ``NonTensorData`` payload survives unwrap.""" + assert unwrap_wire_stripped_payload(NonTensorData(data="hello")) == "hello" + + +# ā”€ā”€ materialize ā€” end-to-end with the wire-stripped tolist shape ā”€ā”€ā”€ā”€ā”€ā”€ + + +def _valid_stack(n: int) -> NonTensorStack: + """A real ``NonTensorStack`` we can patch ``.tolist()`` on. + + Contents are irrelevant ā€” ``materialize`` only iterates the items + returned by ``tolist()``, which we override below. + """ + return NonTensorStack(*(NonTensorData(data=None) for _ in range(n))) + + +def test_materialize_handles_wire_stripped_nontensor_stack() -> None: + """A stack of empty TDs materializes to an object array of ``None``.""" + items = [TensorDict({}, batch_size=[]) for _ in range(4)] + stack = _valid_stack(4) + with patch.object(stack, "tolist", return_value=items): + td = TensorDict({"content": stack}, batch_size=[4]) + bdd = materialize(td, layout="padded") + + arr = bdd["content"] + assert isinstance(arr, np.ndarray) + assert arr.dtype == object + assert arr.shape == (4,) + assert list(arr) == [None, None, None, None] + + +def test_materialize_preserves_real_nontensor_data() -> None: + """Real ``NonTensorStack`` of strings materializes to the raw strings. + + Guards against the wire-stripped fix accidentally substituting + ``None`` for legitimate string content (the happy path that + Mooncake's pickle wire and the patched simple-backend wire produce). + """ + real = NonTensorStack( + NonTensorData(data="hello"), + NonTensorData(data="world"), + NonTensorData(data="!"), + ) + td = TensorDict({"content": real}, batch_size=[3]) + + bdd = materialize(td, layout="padded") + + arr = bdd["content"] + assert isinstance(arr, np.ndarray) + assert arr.dtype == object + assert arr.shape == (3,) + assert list(arr) == ["hello", "world", "!"] + + +def test_materialize_decodes_nontensor_stack_with_tensor_field() -> None: + """Per-field decode: tensor fields stay padded while object fields ride. + + Guards the invariant that ``materialize``'s object-decode is + per-field, not all-or-nothing ā€” a TensorDict can mix jagged tensor + leaves and ``NonTensorStack`` leaves in the same put. + """ + ids_padded = torch.tensor( + [[10, 20, 30, 0], [40, 50, 0, 0], [60, 70, 80, 90]], dtype=torch.long + ) + lens = torch.tensor([3, 2, 4], dtype=torch.long) + ids_nested = to_nested_by_length(ids_padded, lens) + msg = NonTensorStack({"id": 0}, {"id": 1}, {"id": 2}) + + td = TensorDict( + {"input_ids": ids_nested, "message_log": msg}, + batch_size=[3], + ) + + bdd = materialize( + td, + layout="padded", + pad_value_dict={"input_ids": 999}, + ) + + # Tensor field padded with 999 as usual. + assert bdd["input_ids"][1, 2].item() == 999 + # Object field comes back as np.ndarray(object). + assert isinstance(bdd["message_log"], np.ndarray) + assert bdd["message_log"].dtype == object + assert [d["id"] for d in bdd["message_log"]] == [0, 1, 2] + + +# Real production end-to-end coverage of object columns (put ā†’ wire ā†’ +# get ā†’ decode) against both TQ backends lives in +# tests/data_plane/functional/test_tq_lifecycle.py::test_object_round_trip_backends +# and ::test_object_and_tensor_mixed_round_trip_backends. The unit +# tests above cover the decode path in isolation; the functional tests +# cover the full wire round-trip. + + +def test_materialize_realistic_message_log_object_field() -> None: + """Realistic multi-turn message_log decomposes into ``turn_roles`` / + ``turn_contents`` as ``np.ndarray(dtype=object)`` and materializes back.""" + + n = 4 + ml_batch = make_multi_turn_message_log(n=n, turns_per_sample=[1, 2, 3, 4], seed=51) + decomposed = decompose_message_log(ml_batch) + + # The wire-shape: turn_roles + turn_contents are per-sample lists. + # Build a TD with a NonTensorStack of those lists. + roles_stack = NonTensorStack(*[list(r) for r in decomposed["turn_roles"]]) + contents_stack = NonTensorStack(*[list(c) for c in decomposed["turn_contents"]]) + td = TensorDict( + { + "turn_lengths": decomposed["turn_lengths"], + "turn_roles": roles_stack, + "turn_contents": contents_stack, + }, + batch_size=[n], + ) + + out = materialize(td, layout="padded") + # Object fields come back as np.ndarray(dtype=object) ā€” the codec's + # canonical decode of NonTensorStack. + assert isinstance(out["turn_roles"], np.ndarray) + assert out["turn_roles"].dtype == object + assert isinstance(out["turn_contents"], np.ndarray) + # Per-sample identity survives the decode. + for i in range(n): + assert list(out["turn_roles"][i]) == list(decomposed["turn_roles"][i]) diff --git a/tests/unit/data_plane/test_correctness.py b/tests/unit/data_plane/test_correctness.py new file mode 100644 index 00000000000..986e53097aa --- /dev/null +++ b/tests/unit/data_plane/test_correctness.py @@ -0,0 +1,480 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Correctness invariants for the sync 1-hop data-plane. + +Each test guards a real bug we either hit (Mapping check, tensordict +import, clear_samples ordering) or could silently introduce. Tests target +the ABC contract through ``NoOpDataPlaneClient``, so they run without +TQ installed. +""" + +from __future__ import annotations + +import pytest +import torch +from tensordict import TensorDict + +from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient +from nemo_rl.data_plane.column_io import kv_first_write, read_columns, write_columns +from nemo_rl.data_plane.interfaces import KVBatchMeta +from nemo_rl.data_plane.preshard import shard_meta_for_dp +from nemo_rl.data_plane.schema import DP_TRAIN_FIELDS +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + +from ._rollout_shapes import ( + keys_from_uids, + make_rollout_batch, + register_train_partition, +) + +# ā”€ā”€ helpers ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def _final_batch(n: int = 4, *, with_image: bool = False) -> BatchedDataDict: + d: BatchedDataDict = BatchedDataDict() + d["input_ids"] = torch.arange(n * 8, dtype=torch.long).reshape(n, 8) + d["input_lengths"] = torch.tensor([8] * n, dtype=torch.long) + d["token_mask"] = torch.ones((n, 8), dtype=torch.long) + d["sample_mask"] = torch.ones((n,), dtype=torch.long) + d["generation_logprobs"] = torch.zeros((n, 8), dtype=torch.float32) + if with_image: + # Multimodal extras ā€” exercises the "any tensor field" branch + # in kv_first_write. + d["image_features"] = torch.randn((n, 16, 32), dtype=torch.bfloat16) + return d + + +# ā”€ā”€ fail-loud invariants ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_kv_batch_get_after_clear_raises() -> None: + """Real bug guard: v3 driver tried to read input_ids for log_data + AFTER clear_samples, hit ``ValueError: keys not found``. We now stash + before clear ā€” this test pins the contract that get-after-clear + must fail loud, not silently return empty.""" + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=2) + fb = _final_batch(2) + meta = kv_first_write( + fb, + sample_ids=keys_from_uids(["a", "b"]), + dp_client=client, + partition_id="train", + ) + + client.clear_samples(sample_ids=meta.sample_ids, partition_id="train") + + with pytest.raises(KeyError): + # NoOp raises KeyError when the partition entry is gone. + client.get_samples( + sample_ids=meta.sample_ids, + partition_id="train", + select_fields=["input_ids"], + ) + + +def test_kv_batch_get_unproduced_field_raises() -> None: + """Mid-pipeline guard: requesting a field that no producer has + written must fail loud, not return zeros / silently skip.""" + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=2) + fb = _final_batch(2) + meta = kv_first_write( + fb, + sample_ids=keys_from_uids(["a", "b"]), + dp_client=client, + partition_id="train", + ) + + # ``advantages`` has not been written yet (driver delta-write). + with pytest.raises(KeyError): + client.get_samples( + sample_ids=meta.sample_ids, + partition_id="train", + select_fields=["advantages"], + ) + + +def test_get_data_without_select_fields_raises() -> None: + """P2 invariant ā€” never silently fetch all fields.""" + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=2) + fb = _final_batch(2) + kv_first_write( + fb, + sample_ids=keys_from_uids(["a", "b"]), + dp_client=client, + partition_id="train", + ) + + bare_meta = KVBatchMeta( + partition_id="train", + task_name="train", + sample_ids=["a_g0", "b_g0"], + fields=None, # no fields on meta + ) + with pytest.raises(ValueError, match=r"select_fields|fields"): + client.get_data(bare_meta, select_fields=None) + + +def test_kv_batch_put_rejects_non_tensor_leaves() -> None: + """P3 ā€” no pickle on the bus. Adapters MUST reject non-tensor + leaves so callers can't accidentally ship Python objects.""" + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=2, fields=["input_ids", "metadata"]) + + # Build a TensorDict that smuggles a non-tensor ā€” bypass via + # tensordict's NonTensorData where possible. + from tensordict import NonTensorData + + bad_td = TensorDict( + { + "input_ids": torch.zeros((2, 4), dtype=torch.long), + "metadata": NonTensorData(["a", "b"], batch_size=[2]), + }, + batch_size=[2], + ) + with pytest.raises(TypeError, match=r"non-tensor"): + client.put_samples( + sample_ids=["x_g0", "y_g0"], + partition_id="train", + fields=bad_td, + ) + + +def test_claim_meta_unregistered_task_raises() -> None: + """Catches typo'd consumer task names early.""" + client = NoOpDataPlaneClient() + client.register_partition( + partition_id="train", + fields=["input_ids"], + num_samples=2, + consumer_tasks=["lp"], + ) + with pytest.raises(KeyError, match=r"task"): + client.claim_meta( + partition_id="train", + task_name="trian", # typo + required_fields=["input_ids"], + batch_size=2, + ) + + +# ā”€ā”€ lifecycle invariants ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_kv_clear_with_none_drops_partition() -> None: + """Step-end teardown must remove the partition entirely so the + next step's register_partition starts clean.""" + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=2) + fb = _final_batch(2) + meta = kv_first_write( + fb, + sample_ids=keys_from_uids(["a", "b"]), + dp_client=client, + partition_id="train", + ) + + client.clear_samples(sample_ids=None, partition_id="train") + + # Partition is gone ā€” re-registering must succeed. + register_train_partition(client, num_samples=2) + + +def test_double_register_partition_is_idempotent_overwrite() -> None: + """Re-registering the same partition_id within a step (e.g. retry) + must overwrite cleanly, not append fields.""" + client = NoOpDataPlaneClient() + client.register_partition( + partition_id="train", + fields=["a"], + num_samples=2, + consumer_tasks=["t"], + ) + client.register_partition( + partition_id="train", + fields=["b"], + num_samples=4, + consumer_tasks=["t"], + ) + rec = client._partitions["train"] + assert rec.fields == ["b"] + assert rec.num_samples == 4 + + +def test_check_consumption_status_only_true_when_all_consumed() -> None: + """Authoritative cross-worker stage-done signal ā€” must NOT lie + when consumers haven't fetched yet.""" + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=2) + fb = _final_batch(2) + meta = kv_first_write( + fb, + sample_ids=keys_from_uids(["a", "b"]), + dp_client=client, + partition_id="train", + ) + # No consumer has fetched yet. + assert not client.check_consumption_status("train", ["train"]) + + # Simulate the worker fetch. + client.claim_meta( + partition_id="train", + task_name="train", + required_fields=["input_ids"], + batch_size=meta.size, + ) + assert client.check_consumption_status("train", ["train"]) + + +# ā”€ā”€ per-DP shard invariants ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_shard_meta_for_dp_partitions_keys_disjointly() -> None: + """Sum of shard sizes == total, and pairwise disjoint. + + ``shard_meta_for_dp`` returns ``(list[KVBatchMeta], unsorted_indices)``; + here we only care about the metas. + """ + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=8) + fb = _final_batch(8) + meta = kv_first_write( + fb, + sample_ids=keys_from_uids([f"u{i}" for i in range(8)]), + dp_client=client, + partition_id="train", + ) + + shards, _ = shard_meta_for_dp(meta, dp_world=4, batch_size=8) + assert len(shards) == 4 + assert sum(len(s.sample_ids) for s in shards) == len(meta.sample_ids) + seen: set[str] = set() + for s in shards: + for k in s.sample_ids: + assert k not in seen, f"duplicate key {k!r} across DP shards" + seen.add(k) + assert seen == set(meta.sample_ids) + + +def test_shard_meta_for_dp_keeps_partition_id() -> None: + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=4) + fb = _final_batch(4) + meta = kv_first_write( + fb, + sample_ids=keys_from_uids([f"u{i}" for i in range(4)]), + dp_client=client, + partition_id="train", + ) + shards, _ = shard_meta_for_dp(meta, dp_world=2, batch_size=4) + for s in shards: + assert s.partition_id == meta.partition_id + assert s.task_name == meta.task_name + + +# ā”€ā”€ multimodal / VLM extras ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_kv_first_write_carries_multimodal_extras_through_tq() -> None: + """End-to-end flow for VLM: image features must round-trip via TQ + with original shape + dtype, not be silently dropped or coerced.""" + client = NoOpDataPlaneClient() + fields = list(DP_TRAIN_FIELDS) + ["image_features"] + client.register_partition( + partition_id="train", + fields=fields, + num_samples=4, + consumer_tasks=["train"], + ) + fb = _final_batch(4, with_image=True) + expected = fb["image_features"].clone() + + meta = kv_first_write( + fb, + sample_ids=keys_from_uids([f"u{i}" for i in range(4)]), + dp_client=client, + partition_id="train", + ) + assert "image_features" in meta.fields + + fetched = read_columns(client, meta, select_fields=["image_features"]) + got = fetched["image_features"] + assert got.shape == expected.shape + assert got.dtype == expected.dtype, ( + f"dtype drift: expected {expected.dtype}, got {got.dtype}" + ) + assert torch.equal(got, expected) + + +# ā”€ā”€ dtype preservation ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_kv_batch_put_preserves_bf16_dtype() -> None: + """Catches silent fp32 promotion in the put path.""" + client = NoOpDataPlaneClient() + client.register_partition( + partition_id="train", + fields=["x"], + num_samples=2, + consumer_tasks=["train"], + ) + x = torch.randn((2, 4), dtype=torch.bfloat16) + td = TensorDict({"x": x}, batch_size=[2]) + client.put_samples(sample_ids=["a", "b"], partition_id="train", fields=td) + + out = client.get_samples( + sample_ids=["a", "b"], partition_id="train", select_fields=["x"] + ) + assert out["x"].dtype == torch.bfloat16 + + +def test_kv_batch_put_preserves_int64_dtype() -> None: + """input_ids is int64; never coerce to int32 silently.""" + client = NoOpDataPlaneClient() + client.register_partition( + partition_id="train", + fields=["input_ids"], + num_samples=2, + consumer_tasks=["train"], + ) + x = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.long) + td = TensorDict({"input_ids": x}, batch_size=[2]) + client.put_samples(sample_ids=["a", "b"], partition_id="train", fields=td) + + out = client.get_samples( + sample_ids=["a", "b"], + partition_id="train", + select_fields=["input_ids"], + ) + assert out["input_ids"].dtype == torch.long + assert torch.equal(out["input_ids"], x) + + +# ā”€ā”€ BatchedDataDict / Mapping check ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_write_columns_accepts_batched_data_dict_input() -> None: + """Real bug guard (job 11614968 v2 crash): worker write-back + silently skipped because BatchedDataDict inherits from UserDict, + not dict. The fix uses ``isinstance(result, Mapping)``; this test + pins that contract. + """ + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=2) + fb = _final_batch(2) + meta = kv_first_write( + fb, + sample_ids=keys_from_uids(["a", "b"]), + dp_client=client, + partition_id="train", + ) + + bdd = BatchedDataDict() + bdd["advantages"] = torch.full((2,), 3.0) + + # write_columns accepts plain dict; the Mapping-check on the worker + # side ensures BatchedDataDict (UserDict) also goes through. + write_columns(client, meta, dict(bdd)) + + out = read_columns(client, meta, select_fields=["advantages"]) + assert torch.equal(out["advantages"], torch.full((2,), 3.0)) + + +# ā”€ā”€ kv_first_write key-mint contract ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_kv_first_write_rejects_key_count_mismatch() -> None: + """If ``len(keys) != n_samples``, keys would silently mis-align. + Must fail loud. (Caller-side ``n % len(uids) == 0`` is now enforced + at the rollout actor ā€” see ``SyncRolloutActor.rollout_and_first_put``.)""" + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=5) + fb = _final_batch(5) + with pytest.raises(ValueError, match=r"must match batch size"): + kv_first_write( + fb, + sample_ids=["a_g0", "b_g0"], # 2 keys for a 5-sample batch + dp_client=client, + partition_id="train", + ) + + +def test_kv_first_write_meta_sequence_lengths_match_input_lengths() -> None: + """meta.sequence_lengths is consumed by Megatron's balanced packing + on the driver ā€” it MUST mirror final_batch.input_lengths.""" + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=4) + fb = _final_batch(4) + fb["input_lengths"] = torch.tensor([3, 5, 7, 8], dtype=torch.long) + + meta = kv_first_write( + fb, + sample_ids=keys_from_uids([f"u{i}" for i in range(4)]), + dp_client=client, + partition_id="train", + ) + assert meta.sequence_lengths == [3, 5, 7, 8] + + +# ā”€ā”€ Realistic-shape round-trip ā”€ā”€ +# Uses ``_rollout_shapes.make_rollout_batch`` so the put/read path is +# exercised with the same dtypes (bf16 logprobs, int32 masks, int64 ids) +# and realistic value distributions a production rollout produces. + + +def test_kv_first_write_then_read_preserves_dtypes_realistic() -> None: + """Full kv_first_write ā†’ get_samples round-trip preserves every field's dtype.""" + + n = 8 + batch = make_rollout_batch(n=n, max_seqlen=128, seed=99) + client = NoOpDataPlaneClient() + client.register_partition( + partition_id="train", + fields=list(DP_TRAIN_FIELDS), + num_samples=n, + consumer_tasks=["train"], + ) + seed = BatchedDataDict( + { + "input_ids": batch["input_ids"], + "input_lengths": batch["input_lengths"], + "token_mask": batch["token_mask"], + "sample_mask": batch["sample_mask"], + "generation_logprobs": batch["generation_logprobs"], + } + ) + meta = kv_first_write( + seed, + sample_ids=[f"u{i}" for i in range(n)], + dp_client=client, + partition_id="train", + ) + out = read_columns( + client, + meta, + select_fields=[ + "input_ids", + "input_lengths", + "token_mask", + "sample_mask", + "generation_logprobs", + ], + ) + assert out["input_ids"].dtype == torch.long + assert out["token_mask"].dtype == torch.int32 + assert out["generation_logprobs"].dtype == torch.bfloat16 + # Per-row lengths preserved. + assert torch.equal(out["input_lengths"].to(torch.long), batch["input_lengths"]) diff --git a/tests/unit/data_plane/test_factory.py b/tests/unit/data_plane/test_factory.py new file mode 100644 index 00000000000..0fe85abbb85 --- /dev/null +++ b/tests/unit/data_plane/test_factory.py @@ -0,0 +1,65 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Plan Ā§4.3 ā€” production factory rejects disabled and unknown impls. + +NoOp via factory is forbidden by design (plan Ā§4.8 R-C10). The +NoOpDataPlaneClient is reachable only as a direct import from tests ā€” +verified by the architecture invariants in test_architecture_invariants. +""" + +from __future__ import annotations + +import pytest + +from nemo_rl.data_plane import build_data_plane_client + + +def test_factory_none_cfg_rejected(): + """T1-factory-none-cfg ā€” None config must fail-fast, not silently + construct anything.""" + with pytest.raises(ValueError): + build_data_plane_client(None) + + +def test_factory_disabled_rejected(): + """T1-factory-disabled-rejected ā€” production factory must not + silently hand back a NoOp on enabled=False.""" + with pytest.raises(ValueError, match=r"disabled|enabled"): + build_data_plane_client({"enabled": False, "impl": "transfer_queue"}) + + +def test_factory_noop_impl_rejected(): + """T1-factory-noop-rejected-in-prod ā€” NoOp is not selectable from + the factory. Catches R-C10 (NoOp leaks into production).""" + with pytest.raises(ValueError): + build_data_plane_client({"enabled": True, "impl": "noop"}) + + +def test_factory_unknown_impl_rejected(): + """T1-factory-unknown-impl ā€” unknown impl name fails-fast with a + message naming the offending value.""" + with pytest.raises(ValueError, match=r"unknown.*impl"): + build_data_plane_client({"enabled": True, "impl": "no_such_thing"}) + + +def test_factory_disabled_error_message_helpful(): + """When the factory rejects a disabled config, the error message + should point users at the legacy trainer escape hatch.""" + with pytest.raises(ValueError) as excinfo: + build_data_plane_client({"enabled": False, "impl": "transfer_queue"}) + msg = str(excinfo.value) + # Some pointer to the legacy path so users can self-recover. + assert "grpo" in msg.lower() or "legacy" in msg.lower(), ( + f"factory rejection should reference the legacy trainer; got: {msg}" + ) diff --git a/tests/unit/data_plane/test_interface_contract.py b/tests/unit/data_plane/test_interface_contract.py new file mode 100644 index 00000000000..3426c3b5067 --- /dev/null +++ b/tests/unit/data_plane/test_interface_contract.py @@ -0,0 +1,126 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""ABC contract test, parameterized over every adapter. + +Every new adapter (TQ today, ``nv-dataplane`` later) must pass this. The +test runs against the NoOp adapter by default ā€” it doesn't require TQ to +be installed, so CI exercises the contract on every push. +""" + +from __future__ import annotations + +import pytest +import torch +from tensordict import TensorDict + +from nemo_rl.data_plane import ( + DataPlaneClient, + KVBatchMeta, + build_data_plane_client, +) +from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient + + +def _build_noop() -> DataPlaneClient: + return NoOpDataPlaneClient() + + +@pytest.fixture(params=[_build_noop], ids=["noop"]) +def client(request) -> DataPlaneClient: + c = request.param() + yield c + c.close() + + +def test_factory_disabled_raises(): + """Factory has no NoOp fallback ā€” disabled config must not reach it. + The legacy trainer (grpo.grpo_train) never calls the factory at all.""" + with pytest.raises(ValueError): + build_data_plane_client({"enabled": False, "impl": "transfer_queue"}) + + +def test_factory_unknown_impl_raises(): + with pytest.raises(ValueError): + build_data_plane_client({"enabled": True, "impl": "noop"}) + + +def test_register_put_get_clear(client: DataPlaneClient): + client.register_partition( + partition_id="p", fields=["x"], num_samples=4, consumer_tasks=["read"] + ) + keys = ["a", "b", "c", "d"] + fields = TensorDict({"x": torch.arange(4)}, batch_size=[4]) + client.put_samples(sample_ids=keys, partition_id="p", fields=fields) + + out = client.get_samples(sample_ids=keys, partition_id="p", select_fields=["x"]) + assert torch.equal(out["x"], torch.arange(4)) + + client.clear_samples(sample_ids=None, partition_id="p") + with pytest.raises(KeyError): + client.get_samples(sample_ids=keys, partition_id="p", select_fields=["x"]) + + +def test_claim_meta_advances_consumption(client: DataPlaneClient): + client.register_partition( + partition_id="p", + fields=["x"], + num_samples=2, + consumer_tasks=["read"], + ) + fields = TensorDict({"x": torch.tensor([10, 20])}, batch_size=[2]) + client.put_samples(sample_ids=["a", "b"], partition_id="p", fields=fields) + + meta = client.claim_meta( + partition_id="p", task_name="read", required_fields=["x"], batch_size=2 + ) + assert isinstance(meta, KVBatchMeta) + assert meta.size == 2 + assert client.check_consumption_status("p", ["read"]) + + +def test_get_data_requires_field_selection(client: DataPlaneClient): + """P2 ā€” silently fetching all fields is forbidden.""" + client.register_partition( + partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["read"] + ) + client.put_samples( + sample_ids=["a"], + partition_id="p", + fields=TensorDict({"x": torch.tensor([1])}, batch_size=[1]), + ) + bare = KVBatchMeta(partition_id="p", task_name=None, sample_ids=["a"], fields=None) + with pytest.raises(ValueError): + client.get_data(bare) + + +def test_kv_batch_put_rejects_non_tensor_leaves(client: DataPlaneClient): + """P3 ā€” adapter must reject non-tensor leaves in the fields TensorDict. + + Uses ``NonTensorData`` (the supported tensordict primitive for + storing arbitrary Python objects in a TensorDict) ā€” a plain string + in a regular TensorDict construction silently disappears in some + tensordict versions, so we'd never reach the validator. + """ + NonTensorData = pytest.importorskip("tensordict").NonTensorData + client.register_partition( + partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["read"] + ) + bad = TensorDict({"x": NonTensorData("hello")}, batch_size=[1]) + with pytest.raises(TypeError, match=r"non-tensor"): + client.put_samples(sample_ids=["a"], partition_id="p", fields=bad) + + +def test_close_is_idempotent(client: DataPlaneClient): + client.close() + client.close() diff --git a/tests/unit/data_plane/test_kvbatchmeta.py b/tests/unit/data_plane/test_kvbatchmeta.py new file mode 100644 index 00000000000..4774c44f1e3 --- /dev/null +++ b/tests/unit/data_plane/test_kvbatchmeta.py @@ -0,0 +1,183 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Plan Ā§4.4 ā€” KVBatchMeta dataclass invariants and pickle survival. + +Key risk caught here: ``KVBatchMeta`` must survive ``cloudpickle`` round +trips (R-H1) ā€” Ray uses cloudpickle for actor dispatch; if the meta +breaks in transit, every TQ-mediated dispatch raises mid-step. +""" + +from __future__ import annotations + +import pickle + +import pytest + +from nemo_rl.data_plane import KVBatchMeta + +from ._rollout_shapes import make_realistic_tags + + +def test_size_matches_keys(): + """T1-meta-len ā€” ``size`` is the source of truth derived from + ``keys``; the two cannot drift.""" + meta = KVBatchMeta( + partition_id="p", + task_name="t", + sample_ids=["a", "b", "c"], + sequence_lengths=[1, 2, 3], + ) + assert meta.size == 3 + assert meta.size == len(meta.sample_ids) + + +def test_default_fields_and_extra_info_optional(): + """``fields`` and ``sequence_lengths`` default to None; + ``extra_info`` defaults to an empty dict.""" + meta = KVBatchMeta(partition_id="p", task_name="t", sample_ids=[]) + assert meta.fields is None + assert meta.sequence_lengths is None + assert meta.extra_info == {} + + +def test_pickle_roundtrip_structural_equality(): + """T1-meta-cloudpickle-roundtrip ā€” Ray actor dispatch uses + cloudpickle. Use stdlib pickle as a strict subset; if pickle works, + cloudpickle does too.""" + meta = KVBatchMeta( + partition_id="train", + task_name="train", + sample_ids=["k0", "k1", "k2"], + fields=["input_ids", "advantages"], + sequence_lengths=[10, 20, 30], + extra_info={"step": 5}, + ) + rt = pickle.loads(pickle.dumps(meta)) + assert rt.partition_id == meta.partition_id + assert rt.task_name == meta.task_name + assert rt.sample_ids == meta.sample_ids + assert rt.fields == meta.fields + assert rt.sequence_lengths == meta.sequence_lengths + assert rt.extra_info == meta.extra_info + assert rt.size == meta.size + + +def test_keys_with_duplicates_allowed_or_warned(): + """KVBatchMeta does not enforce key uniqueness ā€” that's the + adapter's job (R-H2-style: dup keys at put time should fail). + + This test pins the current behavior: meta accepts any list; dupe + detection is downstream. + """ + meta = KVBatchMeta(partition_id="p", task_name="t", sample_ids=["a", "a"]) + assert meta.size == 2 # no dedup at meta level + + +def test_empty_meta_is_valid(): + """T1-shard-empty-input ā€” an empty meta is a valid value (e.g. a DP + rank with no work after sharding).""" + meta = KVBatchMeta(partition_id="p", task_name="t", sample_ids=[]) + assert meta.size == 0 + # Cloud-pickle survives empty too. + rt = pickle.loads(pickle.dumps(meta)) + assert rt.size == 0 + + +def test_partition_id_is_required(): + """``partition_id`` is positional and required ā€” plan R-M3.""" + with pytest.raises(TypeError): + KVBatchMeta(task_name="t", sample_ids=[]) # type: ignore[call-arg] + + +def test_extra_info_default_is_unique_per_instance(): + """Mutable default trap ā€” two metas should not share the same + ``extra_info`` dict object.""" + a = KVBatchMeta(partition_id="p", task_name="t", sample_ids=[]) + b = KVBatchMeta(partition_id="p", task_name="t", sample_ids=[]) + a.extra_info["x"] = 1 + assert "x" not in b.extra_info + + +def test_tags_align_with_keys(): + """``tags`` must be exactly one dict per key, or ``None``.""" + KVBatchMeta( + partition_id="p", + task_name="t", + sample_ids=["a", "b"], + tags=[{"x": 1}, {"x": 2}], + ) + with pytest.raises(ValueError, match=r"align 1:1"): + KVBatchMeta( + partition_id="p", task_name="t", sample_ids=["a", "b"], tags=[{"x": 1}] + ) + + +def test_tags_travel_with_subset_slice_concat(): + """Per-key tags must follow keys through ``subset`` / ``slice`` / + ``concat`` so consumers can filter on tags without fetching data.""" + m = KVBatchMeta( + partition_id="p", + task_name="t", + sample_ids=["a", "b", "c", "d"], + sequence_lengths=[1, 2, 3, 4], + tags=[{"std": 0.1}, {"std": 0.0}, {"std": 0.3}, {"std": 0.0}], + ) + + survivors = m.subset([0, 2]) + assert survivors.sample_ids == ["a", "c"] + assert survivors.tags == [{"std": 0.1}, {"std": 0.3}] + assert survivors.sequence_lengths == [1, 3] + + front = m.slice(0, 2) + assert front.tags == [{"std": 0.1}, {"std": 0.0}] + + joined = front.concat(m.slice(2, 4)) + assert joined.sample_ids == m.sample_ids + assert joined.tags == m.tags + + +def test_tags_none_when_either_side_missing_in_concat(): + """``concat`` drops tags if either side has none ā€” symmetric with + the ``sequence_lengths`` behavior.""" + with_tags = KVBatchMeta( + partition_id="p", task_name="t", sample_ids=["a"], tags=[{"x": 1}] + ) + without = KVBatchMeta(partition_id="p", task_name="t", sample_ids=["b"]) + assert with_tags.concat(without).tags is None + + +# ā”€ā”€ Realistic tags from the rollout-shapes helper ā”€ā”€ + + +def test_realistic_tags_align_with_keys() -> None: + """Driver-stamped tags (std/total_reward/prompt_id/...) align 1:1 with keys.""" + + n = 16 + sample_ids = [f"u{i}" for i in range(n)] + tags = make_realistic_tags(n, zero_std_fraction=0.25, seed=42) + meta = KVBatchMeta( + partition_id="train", + task_name="train", + sample_ids=sample_ids, + tags=tags, + ) + # Per-row alignment + tag schema preserved. + assert meta.size == n + assert len(meta.tags) == n + for tag in meta.tags: + assert {"std", "total_reward", "prompt_id", "weight_version"} <= set(tag.keys()) + # The zero-std rows are the filter input for dynamic sampling ā€” a realistic + # mix lets the subset/concat logic exercise both branches. + n_zero = sum(1 for t in meta.tags if t["std"] == 0.0) + assert n_zero == n // 4 diff --git a/tests/unit/data_plane/test_leader_broadcast.py b/tests/unit/data_plane/test_leader_broadcast.py new file mode 100644 index 00000000000..5a74f438c42 --- /dev/null +++ b/tests/unit/data_plane/test_leader_broadcast.py @@ -0,0 +1,101 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit test for ``_broadcast_batched_data_dict`` on a 2-rank gloo group. + +Exercises the helper that backs ``_fetch(fetch_policy="leader_broadcast")``. +Runs on CPU (gloo) so it stays in the no-GPU Tier 1 lane. +""" + +from __future__ import annotations + +import os + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from nemo_rl.data_plane.worker_mixin import _broadcast_batched_data_dict +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + +def _worker(rank: int, world_size: int, tmp_init_file: str, q): + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + dist.init_process_group( + backend="gloo", + init_method=f"file://{tmp_init_file}", + rank=rank, + world_size=world_size, + ) + try: + if rank == 0: + data = BatchedDataDict( + { + "input_ids": torch.arange(12, dtype=torch.long).reshape(3, 4), + "input_lengths": torch.tensor([4, 3, 2], dtype=torch.int32), + "scalar_meta": "step_42", + } + ) + else: + data = None + + out = _broadcast_batched_data_dict( + data, is_leader=(rank == 0), src=0, group=dist.group.WORLD + ) + + assert torch.equal( + out["input_ids"], torch.arange(12, dtype=torch.long).reshape(3, 4) + ) + assert torch.equal( + out["input_lengths"], torch.tensor([4, 3, 2], dtype=torch.int32) + ) + assert out["scalar_meta"] == "step_42" + q.put((rank, "ok")) + except Exception as e: # pragma: no cover ā€” surface failures to parent + q.put((rank, f"err: {type(e).__name__}: {e}")) + finally: + dist.destroy_process_group() + + +def test_leader_broadcast_round_trip(tmp_path): + init_file = str(tmp_path / "init") + ctx = mp.get_context("spawn") + q = ctx.Queue() + procs = [ + ctx.Process(target=_worker, args=(rank, 2, init_file, q)) for rank in range(2) + ] + for p in procs: + p.start() + for p in procs: + p.join(timeout=30) + assert p.exitcode == 0, f"worker exited with {p.exitcode}" + + results = sorted([q.get() for _ in range(2)]) + assert results == [(0, "ok"), (1, "ok")], results + + +def test_get_replica_group_default_is_none(): + """TQWorkerMixin._get_replica_group must default to None. + + The base default lets ``_fetch(fetch_policy="leader_broadcast")`` + fall back to the independent path when no backend override exists + (Phase 1 / FSDP2 with TP=CP=PP=1). + """ + from nemo_rl.data_plane.worker_mixin import TQWorkerMixin + + class _Stub(TQWorkerMixin): + pass + + assert _Stub()._get_replica_group() is None diff --git a/tests/unit/data_plane/test_local_node_ip.py b/tests/unit/data_plane/test_local_node_ip.py new file mode 100644 index 00000000000..3c5c1078468 --- /dev/null +++ b/tests/unit/data_plane/test_local_node_ip.py @@ -0,0 +1,176 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for _get_local_node_ip and the MC_TCP_BIND_ADDRESS env-var +assignment in the mooncake_cpu adapter path. + +Covers P3: multi-node correctness of the per-process IP binding. + +The helper rejects two classes of non-routable address: +* link-local (169.254/16, fe80::/10) ā€” APIPA via ``avahi-autoipd`` +* loopback (127.0.0.0/8, ::1) ā€” when ``/etc/hosts`` maps the + hostname to 127.0.0.1 +""" + +from __future__ import annotations + +import os + +import pytest + +# ā”€ā”€ helpers ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def _import_helper(): + """Import _get_local_node_ip from the TQ adapter. + + Returns the function if importable, or None if transfer_queue is absent + (the adapter can't be imported without TQ installed because it calls + socket at module scope only for type annotations ā€” but the function + itself lives in the module-level namespace and only touches socket at + call time, so the import is always safe). + """ + try: + from nemo_rl.data_plane.adapters.transfer_queue import _get_local_node_ip + + return _get_local_node_ip + except ImportError: + return None + + +# ā”€ā”€ tests ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_local_node_ip_skips_link_local(monkeypatch) -> None: + """When gethostbyname returns a link-local address (169.254.x.x), the + helper returns an empty string rather than exposing the non-routable address. + + 169.254.0.0/16 is RFC 3927 APIPA ā€” assigned by avahi-autoipd on usb0 on + this cluster. Announcing that address to Mooncake causes 'connection + refused' on peer nodes. + """ + import socket + + fn = _import_helper() + if fn is None: + pytest.skip("transfer_queue adapter not importable in this environment") + + monkeypatch.setattr(socket, "gethostname", lambda: "fake-host") + monkeypatch.setattr(socket, "gethostbyname", lambda _: "169.254.1.1") + + result = fn() + assert result == "", ( + f"Expected empty string for link-local 169.254.1.1, got {result!r}. " + "Link-local addresses must not be announced to Mooncake peers." + ) + + +def test_local_node_ip_skips_loopback(monkeypatch) -> None: + """When gethostbyname returns the loopback address (127.0.0.1), the + helper returns an empty string rather than announcing an unroutable + address to Mooncake peers. + + Hosts where ``/etc/hosts`` maps the hostname to 127.0.0.1 would + otherwise cause cross-node 'connection refused' on Mooncake. + """ + import socket + + fn = _import_helper() + if fn is None: + pytest.skip("transfer_queue adapter not importable in this environment") + + monkeypatch.setattr(socket, "gethostname", lambda: "fake-host") + monkeypatch.setattr(socket, "gethostbyname", lambda _: "127.0.0.1") + + result = fn() + assert result == "", ( + f"Expected empty string for loopback 127.0.0.1, got {result!r}. " + "Loopback addresses must not be announced to Mooncake peers." + ) + + +def test_local_node_ip_returns_routable(monkeypatch) -> None: + """When gethostbyname returns a routable address, the helper returns it.""" + import socket + + fn = _import_helper() + if fn is None: + pytest.skip("transfer_queue adapter not importable in this environment") + + monkeypatch.setattr(socket, "gethostname", lambda: "fake-host") + monkeypatch.setattr(socket, "gethostbyname", lambda _: "10.65.4.22") + + result = fn() + assert result == "10.65.4.22", ( + f"Expected '10.65.4.22' for a routable address, got {result!r}." + ) + + +def test_local_node_ip_returns_empty_on_exception(monkeypatch) -> None: + """If gethostbyname raises (e.g. DNS not available), the helper returns + an empty string rather than propagating the exception. + + This ensures TQDataPlaneClient.__init__ can still run on nodes with + broken DNS; Mooncake simply won't get a bind hint. + """ + import socket + + fn = _import_helper() + if fn is None: + pytest.skip("transfer_queue adapter not importable in this environment") + + monkeypatch.setattr(socket, "gethostname", lambda: "fake-host") + monkeypatch.setattr( + socket, "gethostbyname", lambda _: (_ for _ in ()).throw(OSError("DNS fail")) + ) + + result = fn() + assert result == "", f"Expected empty string on DNS exception, got {result!r}." + + +def test_mc_tcp_bind_address_overwrites_existing(monkeypatch) -> None: + """TQDataPlaneClient.__init__ uses direct assignment (not os.environ.setdefault) + for MC_TCP_BIND_ADDRESS on the mooncake_cpu path. + + On multi-node runs, Ray actors INHERIT environment variables from the driver + process. If setdefault were used, worker actors on other nodes would keep + the driver's IP, announcing listeners that route back to the head node. + The fix (direct assignment) is verified here: a pre-existing stale value + must be overwritten with the local IP. + """ + import socket + + from nemo_rl.data_plane.adapters.transfer_queue import _get_local_node_ip + + local_ip = "10.65.4.100" + + monkeypatch.setattr(socket, "gethostname", lambda: "worker-node-1") + monkeypatch.setattr(socket, "gethostbyname", lambda _: local_ip) + + # Simulate a stale driver IP inherited via Ray actor env inheritance. + monkeypatch.setenv("MC_TCP_BIND_ADDRESS", "10.65.0.1") + + ip = _get_local_node_ip() + if not ip: + pytest.skip("gethostbyname returned empty in this environment") + + # The adapter's __init__ does: os.environ["MC_TCP_BIND_ADDRESS"] = local_ip + # Replicate that assignment (unit-level; we don't bootstrap a full TQ client). + os.environ["MC_TCP_BIND_ADDRESS"] = ip + + assert os.environ["MC_TCP_BIND_ADDRESS"] == local_ip, ( + f"MC_TCP_BIND_ADDRESS should be {local_ip!r} (this node's IP) " + f"not {os.environ['MC_TCP_BIND_ADDRESS']!r}. " + "Direct assignment is required ā€” setdefault would silently keep the " + "stale driver IP and cause 'connection refused' on peer nodes." + ) diff --git a/tests/unit/data_plane/test_message_log_decompose.py b/tests/unit/data_plane/test_message_log_decompose.py new file mode 100644 index 00000000000..3c448021e8d --- /dev/null +++ b/tests/unit/data_plane/test_message_log_decompose.py @@ -0,0 +1,310 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the ``message_log`` wire-boundary decomposition. + +Sits under ``tests/data_plane/`` rather than ``tests/unit/data/`` so the +heavy ``tests/unit/conftest.py`` (which eagerly imports Ray / the full +nemo_rl model stack) doesn't gate collection. The three helpers under +test are pure-Python and need only ``torch`` / ``numpy`` / +``BatchedDataDict`` at runtime. +""" + +from typing import Any + +import pytest +import torch + +from nemo_rl.data.interfaces import LLMMessageLogType +from nemo_rl.data.llm_message_utils import ( + MESSAGE_LOG_BULK_FIELDS, + attach_message_log_view, + decompose_message_log, + reconstruct_message_log, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + +from ._rollout_shapes import make_multi_turn_message_log + + +def _build_message_log_batch() -> list[LLMMessageLogType]: + return [ + [ + {"role": "user", "content": "Q1", "token_ids": torch.tensor([1, 2, 3])}, + {"role": "assistant", "content": "A1", "token_ids": torch.tensor([4, 5])}, + ], + [ + {"role": "user", "content": "Q2", "token_ids": torch.tensor([6, 7])}, + { + "role": "assistant", + "content": "A2", + "token_ids": torch.tensor([8, 9, 10, 11]), + }, + ], + ] + + +def test_decompose_message_log_basic_shapes() -> None: + out = decompose_message_log(_build_message_log_batch()) + assert out["turn_lengths"].tolist() == [[3, 2], [2, 4]] + assert list(out["turn_roles"][0]) == ["user", "assistant"] + assert list(out["turn_contents"][1]) == ["Q2", "A2"] + # First assistant turn's length per sample. + assert out["response_token_lengths"].tolist() == [2, 4] + + +def test_decompose_message_log_no_assistant_turn() -> None: + out = decompose_message_log( + [[{"role": "user", "content": "U", "token_ids": torch.tensor([1, 2])}]] + ) + assert out["turn_lengths"].tolist() == [[2]] + assert out["response_token_lengths"].tolist() == [0] + + +def test_decompose_message_log_picks_first_assistant() -> None: + """If multiple assistant turns exist, ``response_token_lengths`` takes the first.""" + out = decompose_message_log( + [ + [ + {"role": "user", "content": "U", "token_ids": torch.tensor([1])}, + { + "role": "assistant", + "content": "A1", + "token_ids": torch.tensor([2, 3]), + }, + {"role": "user", "content": "U2", "token_ids": torch.tensor([4])}, + { + "role": "assistant", + "content": "A2", + "token_ids": torch.tensor([5, 6, 7, 8]), + }, + ] + ] + ) + assert out["response_token_lengths"].tolist() == [2] + + +def test_decompose_message_log_jagged_turn_count() -> None: + """Samples with different turn counts pad ``turn_lengths`` with zeros.""" + out = decompose_message_log( + [ + [ + {"role": "user", "content": "U", "token_ids": torch.tensor([1, 2])}, + {"role": "assistant", "content": "A", "token_ids": torch.tensor([3])}, + {"role": "tool", "content": "T", "token_ids": torch.tensor([4, 5, 6])}, + ], + [ + {"role": "user", "content": "U", "token_ids": torch.tensor([7])}, + ], + ] + ) + assert out["turn_lengths"].tolist() == [[2, 1, 3], [1, 0, 0]] + + +def test_decompose_message_log_missing_role_raises() -> None: + """Missing ``role`` surfaces loudly as KeyError rather than producing ``""`` silently.""" + with pytest.raises(KeyError): + decompose_message_log( + [[{"content": "no role here", "token_ids": torch.tensor([1])}]] + ) + + +def test_reconstruct_message_log_roundtrip() -> None: + """decompose ā†’ flatten ā†’ reconstruct returns equivalent message_log.""" + ml_batch = _build_message_log_batch() + decomposed = decompose_message_log(ml_batch) + + flat_per_sample = [torch.cat([m["token_ids"] for m in ml]) for ml in ml_batch] + max_total = max(t.shape[0] for t in flat_per_sample) + input_ids = torch.zeros((len(ml_batch), max_total), dtype=torch.long) + for i, t in enumerate(flat_per_sample): + input_ids[i, : t.shape[0]] = t + + rebuilt = reconstruct_message_log( + input_ids=input_ids, + turn_lengths=decomposed["turn_lengths"], + turn_roles=decomposed["turn_roles"], + turn_contents=decomposed["turn_contents"], + ) + + assert len(rebuilt) == len(ml_batch) + for orig_sample, new_sample in zip(ml_batch, rebuilt): + assert len(orig_sample) == len(new_sample) + for orig_turn, new_turn in zip(orig_sample, new_sample): + assert orig_turn["role"] == new_turn["role"] + assert orig_turn["content"] == new_turn["content"] + assert torch.equal(orig_turn["token_ids"], new_turn["token_ids"]) + + +def test_reconstruct_message_log_returns_views() -> None: + """Per-turn ``token_ids`` must be views into the local ``input_ids`` storage.""" + ml_batch = _build_message_log_batch() + decomposed = decompose_message_log(ml_batch) + input_ids = torch.zeros((2, 6), dtype=torch.long) + input_ids[0, :5] = torch.tensor([1, 2, 3, 4, 5]) + input_ids[1, :6] = torch.tensor([6, 7, 8, 9, 10, 11]) + + rebuilt = reconstruct_message_log( + input_ids=input_ids, + turn_lengths=decomposed["turn_lengths"], + turn_roles=decomposed["turn_roles"], + turn_contents=decomposed["turn_contents"], + ) + + parent_ptr = input_ids.untyped_storage().data_ptr() + for sample in rebuilt: + for turn in sample: + if "token_ids" in turn: + assert turn["token_ids"].untyped_storage().data_ptr() == parent_ptr + + +def test_reconstruct_message_log_attaches_generation_logprobs() -> None: + """``generation_logprobs`` is attached only to assistant turns when provided.""" + ml_batch = _build_message_log_batch() + decomposed = decompose_message_log(ml_batch) + input_ids = torch.zeros((2, 6), dtype=torch.long) + input_ids[0, :5] = torch.tensor([1, 2, 3, 4, 5]) + input_ids[1, :6] = torch.tensor([6, 7, 8, 9, 10, 11]) + gen_logprobs = torch.zeros_like(input_ids, dtype=torch.float32) + + rebuilt = reconstruct_message_log( + input_ids=input_ids, + turn_lengths=decomposed["turn_lengths"], + turn_roles=decomposed["turn_roles"], + turn_contents=decomposed["turn_contents"], + generation_logprobs=gen_logprobs, + ) + + for sample in rebuilt: + for turn in sample: + if turn["role"] == "assistant": + assert "generation_logprobs" in turn + assert turn["generation_logprobs"].shape == turn["token_ids"].shape + else: + assert "generation_logprobs" not in turn + + +def test_attach_message_log_view_populates_batch() -> None: + ml_batch = _build_message_log_batch() + decomposed = decompose_message_log(ml_batch) + input_ids = torch.zeros((2, 6), dtype=torch.long) + input_ids[0, :5] = torch.tensor([1, 2, 3, 4, 5]) + input_ids[1, :6] = torch.tensor([6, 7, 8, 9, 10, 11]) + batch: BatchedDataDict[Any] = BatchedDataDict( + {"input_ids": input_ids, **{k: decomposed[k] for k in MESSAGE_LOG_BULK_FIELDS}} + ) + assert "message_log" not in batch + attach_message_log_view(batch) + assert "message_log" in batch + assert len(batch["message_log"]) == 2 + assert batch["message_log"][0][1]["role"] == "assistant" + + +def test_attach_message_log_view_noop_when_fields_absent() -> None: + """Without decomposed fields, ``attach_message_log_view`` must leave the batch unchanged.""" + batch: BatchedDataDict[Any] = BatchedDataDict({"input_ids": torch.zeros((2, 4))}) + attach_message_log_view(batch) + assert "message_log" not in batch + + +def test_attach_message_log_view_idempotent() -> None: + """Calling twice produces the same shape (no exceptions, no doubled state).""" + ml_batch = _build_message_log_batch() + decomposed = decompose_message_log(ml_batch) + input_ids = torch.zeros((2, 6), dtype=torch.long) + batch: BatchedDataDict[Any] = BatchedDataDict( + {"input_ids": input_ids, **{k: decomposed[k] for k in MESSAGE_LOG_BULK_FIELDS}} + ) + attach_message_log_view(batch) + first_len = len(batch["message_log"]) + attach_message_log_view(batch) + assert len(batch["message_log"]) == first_len + + +# ā”€ā”€ Realistic multi-turn coverage using ``_rollout_shapes.make_multi_turn_message_log`` ā”€ā”€ +# Exercises decompose/reconstruct on the same shape of message_log a real +# multi-turn rollout produces ā€” jagged turn counts (1-4), alternating +# user/assistant roles, variable per-turn token lengths. + + +def test_decompose_realistic_multi_turn_jagged_count() -> None: + """Jagged turn-count message logs (1, 4, 2 turns) round-trip via decompose. + + The realistic shape is what multi-turn rollouts produce ā€” varied + per-sample turn counts. ``decompose_message_log`` must pad shorter + samples' ``turn_lengths`` with zeros without losing role / content + alignment. + """ + + # Force three samples with distinctly different turn counts. + ml_batch = make_multi_turn_message_log(n=3, turns_per_sample=[1, 4, 2], seed=23) + decomposed = decompose_message_log(ml_batch) + + n = len(ml_batch) + max_turns = max(len(s) for s in ml_batch) + + # Shapes + assert decomposed["turn_lengths"].shape == (n, max_turns) + assert len(decomposed["turn_roles"]) == n + assert len(decomposed["turn_contents"]) == n + # Shorter samples' tail turns padded with zero + assert int(decomposed["turn_lengths"][0, 1]) == 0 # 1-turn sample, slot 1 empty + assert int(decomposed["turn_lengths"][2, 2]) == 0 # 2-turn sample, slot 2 empty + # Non-padding positions match the source token counts + for i, sample in enumerate(ml_batch): + for t, turn in enumerate(sample): + assert int(decomposed["turn_lengths"][i, t]) == int( + turn["token_ids"].shape[0] + ) + assert decomposed["turn_roles"][i][t] == turn["role"] + + +def test_decompose_reconstruct_roundtrip_realistic_multi_turn() -> None: + """Full decompose ā†’ reconstruct round-trip on a realistic jagged multi-turn log. + + Existing roundtrip test uses a fixed 2-turn (user/assistant) shape via + ``_build_message_log_batch``. This one exercises the full pipeline on + variable turn counts (1, 3, 4 turns) with alternating roles ā€” the + realistic chat shape the wire actually carries. + """ + + ml_batch = make_multi_turn_message_log(n=3, turns_per_sample=[1, 3, 4], seed=17) + decomposed = decompose_message_log(ml_batch) + + # Build the flat input_ids that the consumer would see on the wire. + flat_per_sample = [torch.cat([m["token_ids"] for m in ml]) for ml in ml_batch] + max_total = max(t.shape[0] for t in flat_per_sample) + input_ids = torch.zeros((len(ml_batch), max_total), dtype=torch.long) + for i, t in enumerate(flat_per_sample): + input_ids[i, : t.shape[0]] = t + + rebuilt = reconstruct_message_log( + input_ids=input_ids, + turn_lengths=decomposed["turn_lengths"], + turn_roles=decomposed["turn_roles"], + turn_contents=decomposed["turn_contents"], + ) + + # Sample-level + turn-level identity through the pipeline. + assert len(rebuilt) == len(ml_batch) + for i, (orig_sample, new_sample) in enumerate(zip(ml_batch, rebuilt)): + assert len(orig_sample) == len(new_sample), ( + f"sample {i}: turn count diverged " + f"orig={len(orig_sample)} != rebuilt={len(new_sample)}" + ) + for t, (orig_turn, new_turn) in enumerate(zip(orig_sample, new_sample)): + assert orig_turn["role"] == new_turn["role"], ( + f"sample {i} turn {t}: role diverged" + ) + assert orig_turn["content"] == new_turn["content"] + assert torch.equal(orig_turn["token_ids"], new_turn["token_ids"]) diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py new file mode 100644 index 00000000000..0d471bc2660 --- /dev/null +++ b/tests/unit/data_plane/test_observability.py @@ -0,0 +1,183 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the lean observability decorator. + +Wraps :class:`NoOpDataPlaneClient` so the tests run in the slim Tier-1 +venv (no TQ, no Ray). The lean shape is one user-injected ``on_event`` +callback plus :meth:`snapshot` for cumulative totals ā€” no ABC, no +built-in sinks. +""" + +from __future__ import annotations + +import pytest +import torch +from tensordict import TensorDict + +from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient +from nemo_rl.data_plane.observability import MetricsDataPlaneClient + +from ._rollout_shapes import make_rollout_batch + + +@pytest.fixture +def wrapped_client(): + events: list[dict] = [] + inner = NoOpDataPlaneClient() + client = MetricsDataPlaneClient(inner, on_event=events.append) + yield client, events + inner.close() + + +def test_put_records_bytes_and_count(wrapped_client): + client, events = wrapped_client + client.register_partition( + partition_id="p", fields=["x"], num_samples=4, consumer_tasks=["read"] + ) + fields = TensorDict({"x": torch.zeros(4, dtype=torch.float32)}, batch_size=[4]) + client.put_samples(sample_ids=["a", "b", "c", "d"], partition_id="p", fields=fields) + + put_events = [e for e in events if e["op"] == "put"] + assert len(put_events) == 1 + e = put_events[0] + assert e["status"] == "ok" + assert e["n_keys"] == 4 + assert e["n_bytes"] == 16 # 4 floats * 4 bytes + assert e["wall_ms"] >= 0 + + +def test_get_records_after_put(wrapped_client): + client, events = wrapped_client + client.register_partition( + partition_id="p", fields=["x"], num_samples=2, consumer_tasks=["read"] + ) + client.put_samples( + sample_ids=["a", "b"], + partition_id="p", + fields=TensorDict({"x": torch.ones(2)}, batch_size=[2]), + ) + out = client.get_samples( + sample_ids=["a", "b"], partition_id="p", select_fields=["x"] + ) + assert torch.equal(out["x"], torch.ones(2)) + + get_events = [e for e in events if e["op"] == "get"] + assert len(get_events) == 1 + assert get_events[0]["n_bytes"] > 0 + + +def test_register_and_clear_recorded(wrapped_client): + client, events = wrapped_client + client.register_partition( + partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["r"] + ) + client.clear_samples(sample_ids=None, partition_id="p") + + ops = [e["op"] for e in events] + assert ops.count("register") == 1 + assert ops.count("clear") == 1 + + +def test_error_status_recorded_and_reraised(wrapped_client): + """Decorator does NOT swallow errors ā€” re-raise after recording.""" + client, events = wrapped_client + with pytest.raises(KeyError): + client.get_samples(sample_ids=["a"], partition_id="nope", select_fields=["x"]) + + err = [e for e in events if e["op"] == "get" and e["status"] == "error"] + assert len(err) == 1 + + +def test_snapshot_accumulates_successful_ops(wrapped_client): + client, _ = wrapped_client + client.register_partition( + partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["r"] + ) + client.put_samples( + sample_ids=["a"], + partition_id="p", + fields=TensorDict({"x": torch.zeros(1)}, batch_size=[1]), + ) + snap = client.snapshot() + assert snap["total_ops"] >= 2 # register + put + assert snap["total_bytes"] >= 4 # 1 float = 4 bytes + + +def test_default_callback_is_noop(): + """Omitting on_event must not raise; the wrapper just forwards.""" + inner = NoOpDataPlaneClient() + client = MetricsDataPlaneClient(inner) + client.register_partition( + partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["r"] + ) + client.close() + + +def test_close_propagates(wrapped_client): + client, _ = wrapped_client + client.close() + # Second close must not raise ā€” NoOp is idempotent. + client.close() + + +def test_factory_wraps_when_observability_enabled(): + """Programmatic wrap path; factory.py uses the same MetricsDataPlaneClient.""" + inner = NoOpDataPlaneClient() + seen: list[dict] = [] + client = MetricsDataPlaneClient(inner, on_event=seen.append) + assert hasattr(client, "snapshot") + client.register_partition( + partition_id="p", fields=["x"], num_samples=1, consumer_tasks=["r"] + ) + assert len(seen) == 1 and seen[0]["op"] == "register" + client.close() + + +def test_observability_records_realistic_rollout_put() -> None: + """Metrics middleware records put-bytes correctly when the put carries a + realistic rollout-shaped batch (bf16 logprobs, int32 masks, int64 ids).""" + + inner = NoOpDataPlaneClient() + seen: list[dict] = [] + client = MetricsDataPlaneClient(inner, on_event=seen.append) + + n = 4 + batch = make_rollout_batch(n=n, max_seqlen=64, seed=71) + client.register_partition( + partition_id="train", + fields=["input_ids", "input_lengths", "generation_logprobs"], + num_samples=n, + consumer_tasks=["train"], + ) + fields = TensorDict( + { + "input_ids": batch["input_ids"], + "input_lengths": batch["input_lengths"], + "generation_logprobs": batch["generation_logprobs"], + }, + batch_size=[n], + ) + client.put_samples( + sample_ids=[f"u{i}" for i in range(n)], + partition_id="train", + fields=fields, + ) + + put_events = [e for e in seen if e["op"] == "put"] + assert len(put_events) == 1 + # Bytes should reflect bf16 logprobs (2 bytes/elem) + int64 ids (8 bytes/elem), + # not a fixed-dtype assumption. Lower bound: at least one full int64 batch. + min_expected = n * 64 * 8 # input_ids alone + assert put_events[0]["n_bytes"] >= min_expected + client.close() diff --git a/tests/unit/data_plane/test_preshard_extras.py b/tests/unit/data_plane/test_preshard_extras.py new file mode 100644 index 00000000000..0c5b9e0d62f --- /dev/null +++ b/tests/unit/data_plane/test_preshard_extras.py @@ -0,0 +1,232 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the rollout first-write helper and the meta-only sharder. + +After the sync 1-hop refactor, ``fan_out_per_rank_metas`` was retired in +favor of: + + * ``kv_first_write`` ā€” single flat ``put_samples`` of every tensor + field in the rollout output (multimodal extras ride along). + * ``shard_meta_for_dp`` ā€” pure key-list split per DP rank, no I/O. + +These tests lock in the schema-extensibility behavior (multimodal +fields propagate) and the meta-sharding contract (no key minting, +identity preserved across shards). +""" + +from __future__ import annotations + +import torch + +from nemo_rl.data_plane import KVBatchMeta +from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient +from nemo_rl.data_plane.column_io import kv_first_write, read_columns +from nemo_rl.data_plane.preshard import shard_meta_for_dp +from nemo_rl.data_plane.schema import DP_TRAIN_FIELDS +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + +from ._rollout_shapes import ( + keys_from_uids, + make_rollout_batch, + register_train_partition, +) + + +def _final_batch(n_samples: int = 4, *, with_extras: bool = False) -> BatchedDataDict: + d: BatchedDataDict = BatchedDataDict() + d["input_ids"] = torch.zeros((n_samples, 8), dtype=torch.long) + d["input_lengths"] = torch.tensor([8] * n_samples, dtype=torch.long) + d["token_mask"] = torch.ones((n_samples, 8), dtype=torch.long) + d["sample_mask"] = torch.ones((n_samples,), dtype=torch.long) + d["generation_logprobs"] = torch.zeros((n_samples, 8), dtype=torch.float32) + if with_extras: + d["pixel_values"] = torch.zeros((n_samples, 3, 4, 4), dtype=torch.float32) + return d + + +# ā”€ā”€ kv_first_write schema extensibility ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_kv_first_write_writes_seed_fields(): + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=4) + fb = _final_batch(4) + uids = [f"u{i}" for i in range(4)] + meta = kv_first_write( + fb, sample_ids=keys_from_uids(uids), dp_client=client, partition_id="train" + ) + # Every tensor field in the input lands in TQ under f"{uid}_g0". + assert meta.sample_ids == [f"u{i}_g0" for i in range(4)] + fetched = client.get_samples( + sample_ids=meta.sample_ids, + partition_id="train", + select_fields=["input_ids", "input_lengths", "token_mask", "sample_mask"], + ) + assert fetched["input_ids"].shape == (4, 8) + + +def test_kv_first_write_carries_multimodal_extras(): + """VLM extras (pixel_values) ride along with no schema declaration.""" + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=4) + fb = _final_batch(4, with_extras=True) + uids = [f"u{i}" for i in range(4)] + meta = kv_first_write( + fb, sample_ids=keys_from_uids(uids), dp_client=client, partition_id="train" + ) + assert "pixel_values" in (meta.fields or []) + fetched = client.get_samples( + sample_ids=meta.sample_ids, + partition_id="train", + select_fields=["pixel_values"], + ) + assert fetched["pixel_values"].shape == (4, 3, 4, 4) + + +def test_kv_first_write_keys_match_uids_x_ngen(): + """Keys round-trip: caller mints ``f"{uid}_g{i}"``, helper preserves them + in ``meta.sample_ids`` byte-for-byte.""" + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=6) + fb = _final_batch(6) # 3 prompts Ɨ 2 generations + uids = ["a", "b", "c"] + keys = keys_from_uids(uids, n_gen=2) + meta = kv_first_write(fb, sample_ids=keys, dp_client=client, partition_id="train") + assert meta.sample_ids == ["a_g0", "a_g1", "b_g0", "b_g1", "c_g0", "c_g1"] + + +# ā”€ā”€ shard_meta_for_dp invariants ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def _meta(n: int) -> KVBatchMeta: + return KVBatchMeta( + partition_id="train", + task_name="train", + sample_ids=[f"k{i}" for i in range(n)], + fields=list(DP_TRAIN_FIELDS), + sequence_lengths=[10 + i for i in range(n)], + extra_info={}, + ) + + +def test_shard_meta_for_dp_partitions_keys_disjointly(): + n, dp = 8, 4 + metas, _ = shard_meta_for_dp(_meta(n), dp_world=dp, batch_size=n) + assert len(metas) == dp + flat = [k for m in metas for k in m.sample_ids] + assert sorted(flat) == sorted(_meta(n).sample_ids) # same set, no dups, no minting + + +def test_shard_meta_for_dp_preserves_partition_id(): + metas, _ = shard_meta_for_dp(_meta(4), dp_world=2, batch_size=4) + assert all(m.partition_id == "train" for m in metas) + + +def test_shard_meta_for_dp_unsorted_round_trip(): + """unsorted_indices must reconstruct the input order from DP-rank concat.""" + n, dp = 8, 4 + metas, unsorted = shard_meta_for_dp(_meta(n), dp_world=dp, batch_size=n) + if unsorted is None: + # No reorder happened ā€” DP-rank concat IS the original order. + return + # Build a tensor whose row i is i; permute via dispatch order; reorder back. + flat = [k for m in metas for k in m.sample_ids] + aggregated = torch.tensor([_meta(n).sample_ids.index(k) for k in flat]) + restored = aggregated[torch.tensor(unsorted)] + assert restored.tolist() == list(range(n)) + + +# ā”€ā”€ meta utility helpers ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_kvbatchmeta_subset_filters_keys_and_seqlens(): + m = _meta(6) + sub = m.subset([1, 3, 5]) + assert sub.sample_ids == ["k1", "k3", "k5"] + assert sub.sequence_lengths == [11, 13, 15] + assert sub.partition_id == m.partition_id + + +def test_kvbatchmeta_concat_joins_keys_and_seqlens(): + m1 = _meta(3) + m2 = _meta(6).subset([3, 4, 5]) + j = m1.concat(m2) + assert j.sample_ids == ["k0", "k1", "k2", "k3", "k4", "k5"] + assert j.sequence_lengths == [10, 11, 12, 13, 14, 15] + + +def test_kvbatchmeta_slice_takes_range(): + m = _meta(5) + s = m.slice(1, 4) + assert s.sample_ids == ["k1", "k2", "k3"] + assert s.sequence_lengths == [11, 12, 13] + + +def test_kvbatchmeta_concat_rejects_partition_mismatch(): + import pytest + + m1 = _meta(2) + m2 = KVBatchMeta( + partition_id="other", + task_name="train", + sample_ids=["x", "y"], + fields=None, + sequence_lengths=[1, 2], + ) + with pytest.raises(ValueError, match=r"partition_ids must match"): + m1.concat(m2) + + +# ā”€ā”€ Realistic multimodal extras via the rollout-shapes helper ā”€ā”€ + + +def test_kv_first_write_realistic_multimodal_round_trip() -> None: + """VLM extras (pixel_values bf16, image_grid_thw int64) flow through + the wire as flat top-level fields and come back intact.""" + + n = 4 + batch = make_rollout_batch(n=n, max_seqlen=64, multimodal=True, seed=33) + client = NoOpDataPlaneClient() + client.register_partition( + partition_id="train", + fields=[ + "input_ids", + "input_lengths", + "sample_mask", + "pixel_values", + "image_grid_thw", + ], + num_samples=n, + consumer_tasks=["train"], + ) + final = BatchedDataDict( + { + "input_ids": batch["input_ids"], + "input_lengths": batch["input_lengths"], + "sample_mask": batch["sample_mask"], + "pixel_values": batch["pixel_values"], + "image_grid_thw": batch["image_grid_thw"], + } + ) + meta = kv_first_write( + final, + sample_ids=[f"u{i}" for i in range(n)], + dp_client=client, + partition_id="train", + ) + out = read_columns(client, meta, select_fields=["pixel_values", "image_grid_thw"]) + # bf16 pixel_values + int64 image_grid_thw survive the wire intact. + assert out["pixel_values"].dtype == torch.bfloat16 + assert out["image_grid_thw"].dtype == torch.long + assert out["pixel_values"].shape[0] == n diff --git a/tests/unit/data_plane/test_seqpack_equivalence.py b/tests/unit/data_plane/test_seqpack_equivalence.py new file mode 100644 index 00000000000..6a508c13558 --- /dev/null +++ b/tests/unit/data_plane/test_seqpack_equivalence.py @@ -0,0 +1,305 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Byte-level equivalence between legacy and TQ seqpack/dynbatch paths. + +Both paths share ``BatchedDataDict.shard_by_batch_size(shards=DP_world, +sequence_packing_args=...)`` for cross-DP balance (Option 1 fix). The only +implementation difference is data transport: legacy hands each shard's +tensors directly to the worker; TQ writes them into the queue, then the +worker reads them back. + +This test isolates the seqpack/dynbatch math from rollout sampling, NCCL +non-determinism, and optimizer steps. If it passes, the only remaining +sources of legacy-vs-TQ run-to-run divergence live outside NeMo-RL. + +Spec: + 1. Build a deterministic ``train_data`` with variable input lengths. + 2. Run ``shard_by_batch_size`` on the driver ā€” this is the *one* call + both paths share. Save its output as the legacy reference. + 3. Round-trip each shard through TQ (``put_samples`` ā†’ + ``get_samples`` ā†’ ``materialize``) and re-attach the per-shard + packing metadata from ``extra_info`` (what + ``train_presharded`` does in production). + 4. Assert each rank's tensors and packing metadata are byte-identical + to the legacy reference. +""" + +from __future__ import annotations + +import pytest +import torch +from tensordict import TensorDict + +pytest.importorskip("ray") +transfer_queue = pytest.importorskip("transfer_queue") # noqa: F841 + +from nemo_rl.data_plane import build_data_plane_client, materialize # noqa: E402 +from nemo_rl.distributed.batched_data_dict import BatchedDataDict # noqa: E402 + +from ._rollout_shapes import mooncake_available + +# Ray is initialized once by the parent autouse fixture +# ``tests/unit/conftest.py::init_ray_cluster`` (mirrors production: NeMo-RL +# inits Ray at startup; the data plane attaches on top). Each test just +# builds a TQ client on the shared Ray and closes it on teardown. + + +# Mirror of the seed-field set in nemo_rl/algorithms/grpo_sync.py. +_DP_SEED_FIELDS = ( + "input_ids", + "input_lengths", + "generation_logprobs", + "prev_logprobs", + "reference_policy_logprobs", + "advantages", + "token_mask", + "sample_mask", +) + +# ā”€ā”€ loud-skip helpers ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + +# ā”€ā”€ fixtures ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def _make_tq_cfg(backend: str) -> dict: + # DataPlaneConfig requires the full schema (see interfaces.py); the + # adapter dereferences ``claim_meta_poll_interval_s`` at construction + # so missing it short-circuits the fixture before any test runs. + # ``global_segment_size`` / ``local_buffer_size`` only matter for + # ``mooncake_cpu`` but are required for schema conformance. + return { + "enabled": True, + "impl": "transfer_queue", + "backend": backend, + "storage_capacity": 1024, + "num_storage_units": 1, + "claim_meta_poll_interval_s": 0.5, + "global_segment_size": 8589934592, # 8 GiB ā€” sized for CI host RAM, not prod + "local_buffer_size": 1073741824, # 1 GiB + } + + +@pytest.fixture( + params=["simple", "mooncake_cpu"], + ids=["simple", "mooncake_cpu"], +) +def tq_client(request): + """Parametrized fixture over simple and mooncake_cpu backends. + + mooncake_cpu is skipped when the mooncake wheel is not installed. + Set NEMO_RL_REQUIRE_MOONCAKE=1 to promote the skip to a loud failure. + + Relies on parent autouse ``init_ray_cluster`` for the Ray runtime. + """ + backend = request.param + if backend == "mooncake_cpu" and not mooncake_available(): + pytest.skip( + "mooncake not installed ā€” skipping mooncake_cpu seqpack equivalence " + "(set NEMO_RL_REQUIRE_MOONCAKE=1 to fail loud)" + ) + client = build_data_plane_client(_make_tq_cfg(backend)) + yield client + client.close() + + +def _make_fake_train_data( + n_samples: int = 64, + max_seqlen: int = 4096, + seed: int = 42, +) -> BatchedDataDict: + """Stand-in for GRPO ``train_data``. + + Variable lengths in ``[256, max_seqlen]`` so the bin packer actually + produces multiple bins per shard ā€” flat-length data would trivially + match. + """ + g = torch.Generator().manual_seed(seed) + input_lengths = torch.randint(256, max_seqlen + 1, (n_samples,), generator=g) + input_ids = torch.zeros((n_samples, max_seqlen), dtype=torch.long) + for i in range(n_samples): + n = int(input_lengths[i]) + input_ids[i, :n] = torch.randint(1, 50000, (n,), generator=g) + return BatchedDataDict( + { + "input_ids": input_ids, + "input_lengths": input_lengths, + "advantages": torch.randn(n_samples, max_seqlen, generator=g), + "token_mask": torch.ones(n_samples, max_seqlen), + "sample_mask": torch.ones(n_samples), + "prev_logprobs": torch.randn(n_samples, max_seqlen, generator=g), + "reference_policy_logprobs": torch.randn( + n_samples, max_seqlen, generator=g + ), + "generation_logprobs": torch.randn(n_samples, max_seqlen, generator=g), + } + ) + + +def _round_trip_shards_through_tq( + tq_client, + pre_shards: list, + partition_id: str, +) -> list[BatchedDataDict]: + """Put each shard's seed fields to TQ, fetch back, attach packing metadata. + + This is the same dance the production driver+worker does: + ``grpo_sync.py`` builds per-rank metas and seeds TQ; ``train_presharded`` + fetches its slice and attaches ``extra_info`` packing metadata. + """ + n_total = sum(int(s["sample_mask"].shape[0]) for s in pre_shards) + tq_client.register_partition( + partition_id=partition_id, + fields=list(_DP_SEED_FIELDS), + num_samples=n_total, + consumer_tasks=["train"], + ) + out: list[BatchedDataDict] = [] + for r, shard in enumerate(pre_shards): + n = int(shard["sample_mask"].shape[0]) + keys = [f"r{r}_s{i}" for i in range(n)] + names = [ + f + for f in _DP_SEED_FIELDS + if f in shard and isinstance(shard[f], torch.Tensor) + ] + fields = TensorDict( + {f: shard[f].detach().contiguous() for f in names}, + batch_size=[n], + ) + tq_client.put_samples( + sample_ids=keys, + partition_id=partition_id, + fields=fields, + ) + td_back = tq_client.get_samples( + sample_ids=keys, + partition_id=partition_id, + select_fields=list(names), + ) + bdd = materialize(td_back, layout="padded") + bdd.micro_batch_indices = shard.micro_batch_indices + bdd.micro_batch_lengths = shard.micro_batch_lengths + bdd.elem_counts_per_gb = shard.elem_counts_per_gb + out.append(bdd) + return out + + +def _assert_shards_byte_equal(legacy, recovered, *, expect_metadata: bool) -> None: + assert len(legacy) == len(recovered), ( + f"shard count mismatch: legacy={len(legacy)} tq={len(recovered)}" + ) + for r, (L, T) in enumerate(zip(legacy, recovered)): + L_tensor_keys = {k for k, v in L.data.items() if isinstance(v, torch.Tensor)} + # TQ only transmits _DP_SEED_FIELDS ā€” non-seed legacy fields are + # out of scope for this test. + common = L_tensor_keys & set(_DP_SEED_FIELDS) + assert common <= set(T.data.keys()), ( + f"rank {r}: TQ shard missing seed fields {common - set(T.data.keys())}" + ) + for k in common: + assert L[k].shape == T[k].shape, ( + f"rank {r} field {k}: shape {L[k].shape} != {T[k].shape}" + ) + assert L[k].dtype == T[k].dtype, ( + f"rank {r} field {k}: dtype {L[k].dtype} != {T[k].dtype}" + ) + assert torch.equal(L[k], T[k]), f"rank {r} field {k}: byte-level mismatch" + if expect_metadata: + assert L.micro_batch_indices == T.micro_batch_indices, ( + f"rank {r} micro_batch_indices mismatch" + ) + assert L.micro_batch_lengths == T.micro_batch_lengths, ( + f"rank {r} micro_batch_lengths mismatch" + ) + assert L.elem_counts_per_gb == T.elem_counts_per_gb, ( + f"rank {r} elem_counts_per_gb mismatch" + ) + + +def test_seqpack_legacy_equals_tq(tq_client): + """Sequence packing: legacy shards == TQ-roundtripped shards (byte-level).""" + DP_WORLD = 4 + GBS = 64 + spa = { + "algorithm": "modified_first_fit_decreasing", + "input_key": "input_ids", + "input_lengths_key": "input_lengths", + "sequence_length_pad_multiple": 64, + "max_tokens_per_microbatch": 4096, + } + data = _make_fake_train_data(n_samples=GBS) + + legacy_shards, _ = data.shard_by_batch_size( + DP_WORLD, + batch_size=GBS, + sequence_packing_args=spa, + ) + tq_pre_shards, _ = data.shard_by_batch_size( + DP_WORLD, + batch_size=GBS, + sequence_packing_args=spa, + ) + recovered = _round_trip_shards_through_tq( + tq_client, + tq_pre_shards, + partition_id="seqpack-eq", + ) + _assert_shards_byte_equal(legacy_shards, recovered, expect_metadata=True) + + +def test_dynbatch_legacy_equals_tq(tq_client): + """Dynamic batching: same equivalence claim as seqpack.""" + DP_WORLD = 4 + GBS = 64 + dba = { + "input_key": "input_ids", + "input_lengths_key": "input_lengths", + "sequence_length_round": 64, + "max_tokens_per_microbatch": 4096, + } + data = _make_fake_train_data(n_samples=GBS) + + legacy_shards, _ = data.shard_by_batch_size( + DP_WORLD, + batch_size=GBS, + dynamic_batching_args=dba, + ) + tq_pre_shards, _ = data.shard_by_batch_size( + DP_WORLD, + batch_size=GBS, + dynamic_batching_args=dba, + ) + recovered = _round_trip_shards_through_tq( + tq_client, + tq_pre_shards, + partition_id="dynbatch-eq", + ) + _assert_shards_byte_equal(legacy_shards, recovered, expect_metadata=True) + + +def test_no_packing_legacy_equals_tq(tq_client): + """Sanity: even without packing/dynbatch the transport should be lossless.""" + DP_WORLD = 4 + GBS = 64 + data = _make_fake_train_data(n_samples=GBS) + + legacy_shards = data.shard_by_batch_size(DP_WORLD, batch_size=GBS) + tq_pre_shards = data.shard_by_batch_size(DP_WORLD, batch_size=GBS) + recovered = _round_trip_shards_through_tq( + tq_client, + tq_pre_shards, + partition_id="nopack-eq", + ) + # No packing ā†’ no micro_batch_* metadata to compare. + _assert_shards_byte_equal(legacy_shards, recovered, expect_metadata=False) diff --git a/tests/unit/data_plane/test_smoke.py b/tests/unit/data_plane/test_smoke.py new file mode 100644 index 00000000000..579abc7bd43 --- /dev/null +++ b/tests/unit/data_plane/test_smoke.py @@ -0,0 +1,121 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tier-0 smoke tests ā€” pre-commit gates. + +Cheapest tier: catches drift in module paths, registry keys, and the +public ABC surface. Each test runs in milliseconds and never touches +real Ray / vLLM / TQ. +""" + +from __future__ import annotations + +import inspect + + +def test_sync_utils_module_imports() -> None: + """Catches FQN drift after the algorithms.sync_utils consolidation.""" + from nemo_rl.experience.sync_rollout_actor import ( + SyncRolloutActor, + kv_first_write, + ) + + # ``SyncRolloutActor`` is wrapped by ``@ray.remote`` into + # ``ActorClass(SyncRolloutActor)`` ā€” the wrapper has no + # ``__name__`` attribute. Check via ``repr`` instead. + assert "SyncRolloutActor" in repr(SyncRolloutActor) + assert callable(kv_first_write) + + +def test_sync_rollout_actor_registered_under_vllm_tier() -> None: + """Multinode runs depend on this ā€” without it, tensordict missing on + worker nodes (real bug seen in job 11614968).""" + from nemo_rl.distributed.ray_actor_environment_registry import ( + get_actor_python_env, + ) + from nemo_rl.distributed.virtual_cluster import PY_EXECUTABLES + + fqn = "nemo_rl.experience.sync_rollout_actor.SyncRolloutActor" + env = get_actor_python_env(fqn) + # Same tier as vLLM workers / AsyncTrajectoryCollector / ReplayBuffer. + # Allow either the resolved exec path or the SYSTEM-override sentinel. + assert env in (PY_EXECUTABLES.VLLM, PY_EXECUTABLES.SYSTEM), ( + f"unexpected env tier for {fqn}: {env!r}" + ) + + +def test_kvbatchmeta_schema_unchanged() -> None: + """Schema break check ā€” KVBatchMeta is the cross-process boundary; + adding/removing a field silently would break adapters that pickle it.""" + from nemo_rl.data_plane.interfaces import KVBatchMeta + + expected_fields = { + "partition_id", + "task_name", + "sample_ids", + "fields", + "sequence_lengths", + "extra_info", + "tags", + } + actual_fields = {f.name for f in KVBatchMeta.__dataclass_fields__.values()} + assert actual_fields == expected_fields, ( + f"KVBatchMeta schema drifted. expected={expected_fields}, " + f"actual={actual_fields}" + ) + + +def test_dataplane_client_abc_surface() -> None: + """Catches accidental ABC method removal / rename ā€” e.g. dropping + ``clear_samples`` would break step-end teardown silently.""" + from nemo_rl.data_plane.interfaces import DataPlaneClient + + expected_methods = { + # task-mediated + "register_partition", + "claim_meta", + "get_data", + "check_consumption_status", + # direct-by-key + "put_samples", + "get_samples", + "clear_samples", + # lifecycle + "close", + } + actual_methods = { + name + for name, member in inspect.getmembers(DataPlaneClient, callable) + if not name.startswith("_") and getattr(member, "__isabstractmethod__", False) + } + assert expected_methods.issubset(actual_methods), ( + f"DataPlaneClient ABC missing methods: {expected_methods - actual_methods}" + ) + + +def test_async_and_sync_actors_share_env_tier() -> None: + """Sync should mirror async's env tier ā€” both drive vLLM and write + tensordict to TQ, so they need the same VLLM venv.""" + from nemo_rl.distributed.ray_actor_environment_registry import ( + get_actor_python_env, + ) + + sync_env = get_actor_python_env( + "nemo_rl.experience.sync_rollout_actor.SyncRolloutActor" + ) + async_env = get_actor_python_env( + "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector" + ) + assert sync_env == async_env, ( + f"Sync vs async env tier drift: sync={sync_env!r}, async={async_env!r}" + ) diff --git a/tests/unit/data_plane/test_sync_one_hop.py b/tests/unit/data_plane/test_sync_one_hop.py new file mode 100644 index 00000000000..51431c9b79c --- /dev/null +++ b/tests/unit/data_plane/test_sync_one_hop.py @@ -0,0 +1,493 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Sync 1-hop unit tests. + +Coverage: + * write_columns / read_columns roundtrip ā€” catches async-without-await + bugs (put_samples returning a coroutine instead of running). The + test that didn't exist when the bug was introduced. + * Per-sample key lifecycle ā€” ``kv_first_write`` mints keys, every + subsequent ``shard_meta_for_dp`` slice references the SAME key set + (verl pattern, no re-minting). + * Slice-only dynamic sampling ā€” filter / cache-merge / overflow-slice + on per-sample tensors plus ``meta.sample_ids``. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from nemo_rl.data_plane import KVBatchMeta +from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient +from nemo_rl.data_plane.column_io import kv_first_write, read_columns, write_columns +from nemo_rl.data_plane.preshard import shard_meta_for_dp +from nemo_rl.data_plane.schema import DP_TRAIN_FIELDS +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + +from ._rollout_shapes import ( + keys_from_uids, + make_realistic_tags, + make_rollout_batch, + register_train_partition, +) + + +def _fake_policy(client): + """Minimal stand-in for ``TQPolicy`` exposing only ``discard_samples``. + + ``_apply_dynamic_sampling`` calls ``policy.discard_samples(uids, partition)`` + to drop filtered rows; we delegate to the noop client's ``clear_samples``. + """ + return SimpleNamespace( + discard_samples=lambda sample_ids, partition_id: client.clear_samples( + sample_ids=sample_ids, partition_id=partition_id + ) + ) + + +def _final_batch(n: int = 4) -> BatchedDataDict: + d: BatchedDataDict = BatchedDataDict() + d["input_ids"] = torch.arange(n * 8, dtype=torch.long).reshape(n, 8) + d["input_lengths"] = torch.tensor([8] * n, dtype=torch.long) + d["token_mask"] = torch.ones((n, 8), dtype=torch.long) + d["sample_mask"] = torch.ones((n,), dtype=torch.long) + d["generation_logprobs"] = torch.zeros((n, 8), dtype=torch.float32) + return d + + +# ā”€ā”€ write_columns / read_columns roundtrip ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ +# +# These tests would have caught the asyncio-without-await bug: +# put_samples used to be an async def; calling it without await +# silently dropped the coroutine. The roundtrip below would have +# returned an empty / stale tensor in that case. + + +def test_write_columns_lands_in_tq(): + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=4) + fb = _final_batch(4) + uids = [f"u{i}" for i in range(4)] + meta = kv_first_write( + fb, sample_ids=keys_from_uids(uids), dp_client=client, partition_id="train" + ) + + # Driver delta-write: simulates advantage compute on the trainer. + delta = {"advantages": torch.full((4,), 7.5)} + write_columns(client, meta, delta) + + fetched = client.get_samples( + sample_ids=meta.sample_ids, + partition_id="train", + select_fields=["advantages"], + ) + assert torch.equal(fetched["advantages"], torch.full((4,), 7.5)) + + +def test_read_columns_returns_only_requested_fields(): + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=4) + fb = _final_batch(4) + uids = [f"u{i}" for i in range(4)] + meta = kv_first_write( + fb, sample_ids=keys_from_uids(uids), dp_client=client, partition_id="train" + ) + + bdd = read_columns(client, meta, ["input_ids", "input_lengths"]) + assert "input_ids" in bdd + assert "input_lengths" in bdd + # token_mask was written but not requested ā€” must not be returned. + assert "token_mask" not in bdd + + +def test_write_then_read_roundtrip_after_train_window(): + """Full lifecycle: rollout puts ā†’ driver delta-writes ā†’ read deltas back.""" + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=4) + fb = _final_batch(4) + uids = [f"u{i}" for i in range(4)] + meta = kv_first_write( + fb, sample_ids=keys_from_uids(uids), dp_client=client, partition_id="train" + ) + + # Simulate the full sync 1-hop trainer-step writes: + write_columns( + client, + meta, + { + "prev_logprobs": torch.full((4, 8), 0.1), + "reference_policy_logprobs": torch.full((4, 8), 0.2), + "advantages": torch.full((4,), 0.3), + }, + ) + + # train_presharded would fetch the union ā€” verify all columns present. + fetched = read_columns( + client, + meta, + [ + "input_ids", + "input_lengths", + "prev_logprobs", + "reference_policy_logprobs", + "advantages", + ], + ) + assert torch.allclose(fetched["prev_logprobs"], torch.full((4, 8), 0.1)) + assert torch.allclose(fetched["reference_policy_logprobs"], torch.full((4, 8), 0.2)) + assert torch.allclose(fetched["advantages"], torch.full((4,), 0.3)) + + +# ā”€ā”€ Per-sample key lifecycle invariant ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +def test_meta_keys_identity_across_dp_shards(): + """``shard_meta_for_dp`` must NOT mint new keys ā€” every per-rank + slice references a subset of the original ``meta.sample_ids``.""" + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=8) + fb = _final_batch(8) + uids = [f"u{i}" for i in range(8)] + meta = kv_first_write( + fb, sample_ids=keys_from_uids(uids), dp_client=client, partition_id="train" + ) + + rank_metas, _ = shard_meta_for_dp(meta, dp_world=4, batch_size=8) + flat = {k for m in rank_metas for k in m.sample_ids} + assert flat == set(meta.sample_ids), ( + "shard_meta_for_dp introduced or dropped keys ā€” should be a " + "pure permutation of the original meta.sample_ids." + ) + # Every rank slice points at the same partition. + assert all(m.partition_id == meta.partition_id for m in rank_metas) + + +def test_kv_clear_uses_meta_keys_minted_at_rollout(): + """The keys cleared at step end are the SAME keys the rollout + actor minted ā€” no minting at any stage in between.""" + client = NoOpDataPlaneClient() + register_train_partition(client, num_samples=4) + fb = _final_batch(4) + uids = [f"u{i}" for i in range(4)] + meta = kv_first_write( + fb, sample_ids=keys_from_uids(uids), dp_client=client, partition_id="train" + ) + rollout_keys = list(meta.sample_ids) + + # Workers / driver write deltas ā€” keys still meta.sample_ids. + write_columns(client, meta, {"advantages": torch.zeros(4)}) + rank_metas, _ = shard_meta_for_dp(meta, dp_world=2, batch_size=4) + for rm in rank_metas: + for k in rm.sample_ids: + assert k in set(rollout_keys), ( + "Rank meta references a key not in the original rollout set" + ) + + client.clear_samples(sample_ids=meta.sample_ids, partition_id="train") + # Cleared keys should no longer fetch. + import pytest + + with pytest.raises(KeyError): + client.get_samples( + sample_ids=meta.sample_ids, + partition_id="train", + select_fields=["input_ids"], + ) + + +# ā”€ā”€ Slice-only dynamic sampling logic ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ +# +# These exercise the private ``_apply_dynamic_sampling`` helper in +# grpo_sync.py without requiring a full trainer to spin up. + + +def _make_driver_carry(rewards: list[float], stds: list[float]) -> BatchedDataDict: + n = len(rewards) + return BatchedDataDict( + { + "total_reward": torch.tensor(rewards, dtype=torch.float32), + "std": torch.tensor(stds, dtype=torch.float32), + "baseline": torch.zeros(n), + "input_lengths": torch.tensor([8] * n, dtype=torch.long), + "loss_multiplier": torch.ones(n), + "truncated": torch.zeros(n, dtype=torch.bool), + "length": torch.tensor([8] * n, dtype=torch.long), + "prompt_ids_for_adv": torch.zeros(n, 4, dtype=torch.long), + } + ) + + +def _seed_meta(client: NoOpDataPlaneClient, prefix: str, n: int) -> KVBatchMeta: + """Stage n keys in TQ so clear_samples has something to remove.""" + register_train_partition(client, num_samples=n) + fb = _final_batch(n) + uids = [f"{prefix}{i}" for i in range(n)] + return kv_first_write( + fb, sample_ids=keys_from_uids(uids), dp_client=client, partition_id="train" + ) + + +def _stamp_filter_tags(meta: KVBatchMeta, stds: list[float]) -> KVBatchMeta: + """Mirror the driver's post-baseline/std step: stamp ``std`` into + ``meta.tags`` so ``_apply_dynamic_sampling`` can read the filter + criterion from the meta alone.""" + meta.tags = [{"std": float(s)} for s in stds] + return meta + + +def test_apply_dynamic_sampling_filters_zero_std(): + """Drops uids whose std == 0 and clears their TQ payload.""" + from nemo_rl.algorithms.grpo_sync import _apply_dynamic_sampling + + client = NoOpDataPlaneClient() + meta = _seed_meta(client, "u", n=4) + _stamp_filter_tags(meta, [0.5, 0.0, 0.5, 0.0]) + sd = _make_driver_carry([1.0, 2.0, 3.0, 4.0], [0.5, 0.0, 0.5, 0.0]) + + pm, ps, pur, complete, ds_metrics, _ = _apply_dynamic_sampling( + meta=meta, + driver_carry=sd, + pending_meta=None, + pending_carry=None, + pending_unfiltered_rewards=[], + train_prompts_size=4, + num_gen_batches=1, + max_gen_batches=10, + policy=_fake_policy(client), + ) + # Only 2 survivors ā†’ not complete (need 4). + assert complete is False + assert pm is not None and len(pm.sample_ids) == 2 + # Surviving uids' total_reward is 1.0 and 3.0 (kept indices [0, 2]). + assert torch.equal(ps["total_reward"], torch.tensor([1.0, 3.0])) + assert ps["filtered_reward"] is ps["total_reward"] or torch.equal( + ps["filtered_reward"], ps["total_reward"] + ) + + # Dropped uids' TQ payload was cleared. + import pytest + + with pytest.raises(KeyError): + client.get_samples( + sample_ids=[meta.sample_ids[1]], + partition_id="train", + select_fields=["input_ids"], + ) + # Surviving uids' payload is still alive. + survivors = client.get_samples( + sample_ids=[meta.sample_ids[0], meta.sample_ids[2]], + partition_id="train", + select_fields=["input_ids"], + ) + assert survivors["input_ids"].shape == (2, 8) + + +def test_apply_dynamic_sampling_completes_when_train_size_reached(): + """When pending cache reaches train_prompts_size, returns complete.""" + from nemo_rl.algorithms.grpo_sync import _apply_dynamic_sampling + + client = NoOpDataPlaneClient() + meta = _seed_meta(client, "u", n=4) + _stamp_filter_tags(meta, [0.5, 0.5, 0.5, 0.5]) + sd = _make_driver_carry([1.0, 2.0, 3.0, 4.0], [0.5, 0.5, 0.5, 0.5]) + + pm, ps, _, complete, ds_metrics, unfiltered = _apply_dynamic_sampling( + meta=meta, + driver_carry=sd, + pending_meta=None, + pending_carry=None, + pending_unfiltered_rewards=[], + train_prompts_size=4, + num_gen_batches=1, + max_gen_batches=10, + policy=_fake_policy(client), + ) + assert complete is True + assert pm is not None and len(pm.sample_ids) == 4 + assert ds_metrics["dynamic_sampling_num_gen_batches"] == 1 + # Unfiltered rewards mirror the input (no filtering happened). + assert torch.equal(unfiltered, torch.tensor([1.0, 2.0, 3.0, 4.0])) + + +def test_apply_dynamic_sampling_overflow_slices_and_clears(): + """When the cache exceeds train_prompts_size, slice + clear_samples discards.""" + from nemo_rl.algorithms.grpo_sync import _apply_dynamic_sampling + + client = NoOpDataPlaneClient() + meta = _seed_meta(client, "u", n=6) + _stamp_filter_tags(meta, [0.5] * 6) + sd = _make_driver_carry([1.0] * 6, [0.5] * 6) + + pm, ps, _, complete, ds_metrics, _ = _apply_dynamic_sampling( + meta=meta, + driver_carry=sd, + pending_meta=None, + pending_carry=None, + pending_unfiltered_rewards=[], + train_prompts_size=4, # only need 4; 2 should be discarded + num_gen_batches=1, + max_gen_batches=10, + policy=_fake_policy(client), + ) + assert complete is True + assert len(pm.sample_ids) == 4 + assert ds_metrics.get("dynamic_sampling_num_discarded_valid_samples") == 2 + # Discarded uids (last 2) cleared from TQ. + import pytest + + with pytest.raises(KeyError): + client.get_samples( + sample_ids=[meta.sample_ids[4]], + partition_id="train", + select_fields=["input_ids"], + ) + + +def test_apply_dynamic_sampling_raises_on_max_gen_batches(): + """Exceeding dynamic_sampling_max_gen_batches must raise loudly.""" + from nemo_rl.algorithms.grpo_sync import _apply_dynamic_sampling + + client = NoOpDataPlaneClient() + meta = _seed_meta(client, "u", n=2) + _stamp_filter_tags(meta, [0.0, 0.0]) + sd = _make_driver_carry([1.0, 2.0], [0.0, 0.0]) # all dropped + + import pytest + + with pytest.raises(ValueError, match=r"max_gen_batches"): + _apply_dynamic_sampling( + meta=meta, + driver_carry=sd, + pending_meta=None, + pending_carry=None, + pending_unfiltered_rewards=[], + train_prompts_size=4, + num_gen_batches=11, + max_gen_batches=10, # exceeded + policy=_fake_policy(client), + ) + + +# ā”€ā”€ Multi-stage TQ lifecycle on a realistic batch ā”€ā”€ +# Walks the same sequence the production sync trainer runs: +# 1. register_partition ā†’ 2. kv_first_write (seed) ā†’ 3. stamp filter tags +# ā†’ 4. worker logprob delta-writes ā†’ 5. driver advantage delta-write +# ā†’ 6. full read of train fields ā†’ 7. clear_samples at step-end. +# Each stage uses data shaped like the real rollout writer's output +# (bf16 logprobs, int64 ids, int32 masks, realistic value distributions). + + +def test_full_sync_step_lifecycle_on_realistic_batch() -> None: + """End-to-end TQ lifecycle test mirroring grpo_train_sync's per-step flow.""" + + _PARTITION = "train" + client = NoOpDataPlaneClient() + n = 8 + max_seqlen = 128 + + # ā”€ā”€ Stage 1: register partition with the schema rollout will write ā”€ā”€ + client.register_partition( + partition_id=_PARTITION, + fields=list(DP_TRAIN_FIELDS), + num_samples=n, + consumer_tasks=["prev_lp", "ref_lp", "train"], + ) + + # ā”€ā”€ Stage 2: rollout writes seed fields via kv_first_write ā”€ā”€ + batch = make_rollout_batch(n=n, max_seqlen=max_seqlen, seed=101) + uids = [f"u{i}" for i in range(n)] + seed_fields = { + "input_ids": batch["input_ids"], + "input_lengths": batch["input_lengths"], + "token_mask": batch["token_mask"], + "sample_mask": batch["sample_mask"], + "generation_logprobs": batch["generation_logprobs"], + } + final = BatchedDataDict(seed_fields) + meta = kv_first_write( + final, + sample_ids=keys_from_uids(uids), + dp_client=client, + partition_id=_PARTITION, + ) + # Sanity: meta carries the per-row lengths the driver needs for packing. + assert meta.sequence_lengths is not None + assert len(meta.sample_ids) == n + # Bf16 logprob survives the put. + seeded = client.get_samples( + sample_ids=meta.sample_ids, + partition_id=_PARTITION, + select_fields=["generation_logprobs"], + ) + assert seeded["generation_logprobs"].dtype == torch.bfloat16 + + # ā”€ā”€ Stage 3: driver stamps per-row tags (filter input for dyn sampling) ā”€ā”€ + tags = make_realistic_tags(n, zero_std_fraction=0.25, seed=101) + meta.tags = tags + assert sum(1 for t in tags if t["std"] == 0.0) == n // 4 + + # ā”€ā”€ Stage 4: workers compute logprob deltas, write back ā”€ā”€ + write_columns( + client, + meta, + fields={ + "prev_logprobs": batch["prev_logprobs"], + "reference_policy_logprobs": batch["reference_policy_logprobs"], + }, + ) + + # ā”€ā”€ Stage 5: driver computes advantages, writes back ā”€ā”€ + write_columns( + client, + meta, + fields={"advantages": batch["advantages"]}, + ) + + # ā”€ā”€ Stage 6: full read of train fields (what train_presharded does) ā”€ā”€ + full = read_columns( + client, + meta, + select_fields=[ + "input_ids", + "input_lengths", + "token_mask", + "sample_mask", + "generation_logprobs", + "prev_logprobs", + "reference_policy_logprobs", + "advantages", + ], + ) + # All fields present, dtypes preserved end-to-end. + assert full["input_ids"].dtype == torch.long + assert full["token_mask"].dtype == torch.int32 + assert full["generation_logprobs"].dtype == torch.bfloat16 + assert full["prev_logprobs"].dtype == torch.bfloat16 + assert full["reference_policy_logprobs"].dtype == torch.bfloat16 + assert full["advantages"].dtype == torch.bfloat16 + # Row count survives the full pipeline. + assert full["input_ids"].shape[0] == n + + # ā”€ā”€ Stage 7: step-end clear (mirror of finish_step) ā”€ā”€ + client.clear_samples(sample_ids=meta.sample_ids, partition_id=_PARTITION) + # Subsequent get must fail loud ā€” the keys are gone. + with pytest.raises(KeyError): + client.get_samples( + sample_ids=[meta.sample_ids[0]], + partition_id=_PARTITION, + select_fields=["input_ids"], + ) diff --git a/tests/unit/data_plane/test_tq_lifecycle.py b/tests/unit/data_plane/test_tq_lifecycle.py new file mode 100644 index 00000000000..6c3da9de120 --- /dev/null +++ b/tests/unit/data_plane/test_tq_lifecycle.py @@ -0,0 +1,405 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Single-node TQ smoke ā€” Stage 1 acceptance. + +Mirrors the recipe in the integration plan Ā§3 / Stage 1: +register ā†’ put ā†’ claim_meta ā†’ get_data ā†’ check_consumption ā†’ clear. + +Skipped when the ``transfer_queue`` package is not installed so CI without +the data-plane extra still passes. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +from tensordict import TensorDict + +transfer_queue = pytest.importorskip("transfer_queue") # noqa: F841 + +from nemo_rl.data_plane import build_data_plane_client +from nemo_rl.data_plane.column_io import kv_first_write, read_columns +from nemo_rl.data_plane.interfaces import KVBatchMeta +from nemo_rl.data_plane.schema import DP_TRAIN_FIELDS +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + +from ._rollout_shapes import mooncake_available + +# ā”€ā”€ loud-skip helpers ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + +# ā”€ā”€ fixtures ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + +@pytest.fixture +def tq_client(): + import ray + + if not ray.is_initialized(): + ray.init(local_mode=False, include_dashboard=False) + + client = build_data_plane_client( + { + "enabled": True, + "impl": "transfer_queue", + "backend": "simple", + "storage_capacity": 1024, + "num_storage_units": 1, + "claim_meta_poll_interval_s": 0.5, + "global_segment_size": 8589934592, # 8 GiB (only read by mooncake_cpu) + "local_buffer_size": 1073741824, # 1 GiB (only read by mooncake_cpu) + } + ) + yield client + client.close() + + +@pytest.fixture( + params=["simple", "mooncake_cpu"], + ids=["simple", "mooncake_cpu"], +) +def tq_client_backends(request): + """Parametrized fixture over simple and mooncake_cpu backends. + + mooncake_cpu is skipped when the mooncake wheel is not installed. + Set NEMO_RL_REQUIRE_MOONCAKE=1 to promote the skip to a loud failure. + """ + backend = request.param + if backend == "mooncake_cpu" and not mooncake_available(): + pytest.skip( + "mooncake not installed ā€” skipping mooncake_cpu backend " + "(set NEMO_RL_REQUIRE_MOONCAKE=1 to fail loud)" + ) + + import ray + + if not ray.is_initialized(): + ray.init(local_mode=False, include_dashboard=False) + + client = build_data_plane_client( + { + "enabled": True, + "impl": "transfer_queue", + "backend": backend, + "storage_capacity": 1024, + "num_storage_units": 1, + "claim_meta_poll_interval_s": 0.5, + "global_segment_size": 8589934592, # 8 GiB + "local_buffer_size": 1073741824, # 1 GiB + } + ) + yield client + client.close() + + +def test_smoke_round_trip(tq_client) -> None: + tq_client.register_partition( + partition_id="smoke", + fields=["x"], + num_samples=4, + consumer_tasks=["read"], + ) + keys = ["a", "b", "c", "d"] + tq_client.put_samples( + sample_ids=keys, + partition_id="smoke", + fields=TensorDict({"x": torch.arange(4)}, batch_size=[4]), + ) + + meta = tq_client.claim_meta( + partition_id="smoke", + task_name="read", + required_fields=["x"], + batch_size=4, + timeout_s=30.0, + ) + assert meta.size == 4 + + data = tq_client.get_data(meta) + # Order may differ from input ā€” match against the meta's keys. + expected = torch.tensor([keys.index(k) for k in meta.sample_ids]) + assert torch.equal(data["x"], expected) + + assert tq_client.check_consumption_status("smoke", ["read"]) + + tq_client.clear_samples(sample_ids=None, partition_id="smoke") + + +def test_smoke_round_trip_backends(tq_client_backends) -> None: + """Smoke round-trip parameterized over both backends. + + Covers P5 (T2-backend-bytewise-equal) ā€” the same put/get lifecycle must + work on simple and mooncake_cpu. mooncake_cpu is skipped when unavailable. + """ + client = tq_client_backends + client.register_partition( + partition_id="smoke-backend", + fields=["x"], + num_samples=4, + consumer_tasks=["read"], + ) + keys = ["a", "b", "c", "d"] + client.put_samples( + sample_ids=keys, + partition_id="smoke-backend", + fields=TensorDict({"x": torch.arange(4)}, batch_size=[4]), + ) + + meta = client.claim_meta( + partition_id="smoke-backend", + task_name="read", + required_fields=["x"], + batch_size=4, + timeout_s=30.0, + ) + assert meta.size == 4 + + data = client.get_data(meta) + expected = torch.tensor([keys.index(k) for k in meta.sample_ids]) + assert torch.equal(data["x"], expected) + + client.clear_samples(sample_ids=None, partition_id="smoke-backend") + + +def test_smoke_round_trip_1d_fields(tq_client) -> None: + """A 1D (N,) tensor put into TQ must come back as (N,), not (N,1). + + Regression guard for R-C2: TQ's KVStorageManager path silently unsqueezes + 1D fields. The adapter's `_promote_1d_leaves` + `_from_wire` pair fix + this for the mooncake_cpu backend; this test verifies simple backend does + not introduce the regression. + """ + n = 6 + reward = torch.arange(n, dtype=torch.float32) + + tq_client.register_partition( + partition_id="smoke-1d", + fields=["reward"], + num_samples=n, + consumer_tasks=["read"], + ) + keys = [f"k{i}" for i in range(n)] + tq_client.put_samples( + sample_ids=keys, + partition_id="smoke-1d", + fields=TensorDict({"reward": reward}, batch_size=[n]), + ) + + meta = tq_client.claim_meta( + partition_id="smoke-1d", + task_name="read", + required_fields=["reward"], + batch_size=n, + timeout_s=30.0, + ) + data = tq_client.get_data(meta) + + assert data["reward"].shape == reward.shape, ( + f"Expected shape {tuple(reward.shape)} for 1D field, " + f"got {tuple(data['reward'].shape)}. " + "TQ must not unsqueeze 1D tensors silently (R-C2)." + ) + + tq_client.clear_samples(sample_ids=None, partition_id="smoke-1d") + + +# ā”€ā”€ Object-field round-trip across backends ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ +# +# Closes the coverage gap: prior tests exercised np.ndarray(object) only via +# the in-process codec (test_codec_object.py) or sent tensor-only fields +# through both backends (test_smoke_round_trip_backends). Sending object +# fields through mooncake_cpu was untested. This test covers that path. + + +def _object_payload(n: int) -> np.ndarray: + """Heterogeneous per-row Python objects, mimicking message_log shape.""" + rows = [ + { + "id": i, + "text": f"sample {i} content " * (i % 5 + 1), # variable-length strings + "tags": [f"t{i}", f"t{i + 1}"], + } + for i in range(n) + ] + arr = np.empty(n, dtype=object) + for i, r in enumerate(rows): + arr[i] = r + return arr + + +def test_object_round_trip_backends(tq_client_backends) -> None: + """np.ndarray(dtype=object) put ā†’ get ā†’ decode equality, both backends. + + Mirrors the wire used by ``SyncRolloutActor.kv_first_write`` for + ``message_log`` / ``content``: object fields ride as + ``np.ndarray(dtype=object)`` (matching ``sync_rollout_actor.py`` + line 273 / 292); the TensorDict constructor wraps them as + ``NonTensorData`` internally. :func:`read_columns` ā†’ + :func:`materialize` decodes them back to ``np.ndarray(dtype=object)``. + """ + client = tq_client_backends + n = 8 + field_name = "msg_log" + keys = [f"obj_{i}" for i in range(n)] + + client.register_partition( + partition_id="obj-backend", + fields=[field_name], + num_samples=n, + consumer_tasks=["read"], + ) + client.put_samples( + sample_ids=keys, + partition_id="obj-backend", + fields=TensorDict( + {field_name: _object_payload(n)}, + batch_size=[n], + ), + ) + meta = KVBatchMeta( + partition_id="obj-backend", + task_name="read", + sample_ids=keys, + fields=[field_name], + ) + + bdd = read_columns(client, meta, select_fields=[field_name]) + + assert isinstance(bdd[field_name], np.ndarray) + assert bdd[field_name].dtype == object + assert bdd[field_name].shape == (n,) + expected = _object_payload(n) + for i in range(n): + assert bdd[field_name][i] == expected[i], ( + f"row {i} mismatch: got {bdd[field_name][i]!r}, expected {expected[i]!r}" + ) + + client.clear_samples(sample_ids=None, partition_id="obj-backend") + + +def test_object_and_tensor_mixed_round_trip_backends(tq_client_backends) -> None: + """End-to-end mirror of ``SyncRolloutActor.kv_first_write``. + + Pins the production e2e GRPO pipeline shape on both backends: + + * ``register_partition`` declares ``DP_TRAIN_FIELDS`` (tensor-only), + matching :meth:`TQPolicy.prepare_step`. + * ``bulk_batch`` includes 1D + 2D tensors **and** an + ``np.ndarray(dtype=object)`` (``content``) ā€” the shape built by + ``sync_rollout_actor.py`` lines 257ā€“273. + * ``kv_first_write`` does the put through :func:`pack_jagged_fields`. + * ``read_columns`` fetches a mixed tensor + object subset, the same + pattern used by ``grpo_sync.py`` lines 887ā€“896. + + Regression guard for the data-plane wire round-trip end-to-end. + """ + client = tq_client_backends + n = 6 + seq_len = 4 + sample_ids = [f"sample_{i}" for i in range(n)] + partition_id = "mix-e2e" + + # Tensor-only schema ā€” matches `TQPolicy.prepare_step`. + client.register_partition( + partition_id=partition_id, + fields=list(DP_TRAIN_FIELDS), + num_samples=n, + consumer_tasks=["read"], + ) + + # Production-shape `bulk_batch`: tensors + np.ndarray(dtype=object). + input_ids = torch.arange(n * seq_len, dtype=torch.long).reshape(n, seq_len) + input_lengths = torch.full((n,), seq_len, dtype=torch.long) + generation_logprobs = torch.zeros(n, seq_len, dtype=torch.float) + token_mask = torch.ones(n, seq_len, dtype=torch.float) + sample_mask = torch.ones(n, dtype=torch.float) + content = _object_payload(n) + + bulk_batch = BatchedDataDict( + { + "input_ids": input_ids, + "input_lengths": input_lengths, + "generation_logprobs": generation_logprobs, + "token_mask": token_mask, + "sample_mask": sample_mask, + "content": content, + } + ) + + # Production write path. + meta = kv_first_write( + bulk_batch, + sample_ids=sample_ids, + dp_client=client, + partition_id=partition_id, + task_name="read", + ) + + # Production read path ā€” mixed tensor + object subset. + bdd = read_columns( + client, meta, select_fields=["input_ids", "input_lengths", "content"] + ) + assert torch.equal(bdd["input_ids"], input_ids) + assert torch.equal(bdd["input_lengths"], input_lengths) + expected = _object_payload(n) + for i in range(n): + assert bdd["content"][i] == expected[i], ( + f"row {i} content mismatch: got {bdd['content'][i]!r}, " + f"expected {expected[i]!r}" + ) + + # Tensor-only subset still works. + only_ids = read_columns(client, meta, select_fields=["input_ids"]) + assert torch.equal(only_ids["input_ids"], input_ids) + assert "content" not in only_ids + + # Object-only subset still works. + only_content = read_columns(client, meta, select_fields=["content"]) + assert isinstance(only_content["content"], np.ndarray) + assert "input_ids" not in only_content + + client.clear_samples(sample_ids=None, partition_id=partition_id) + + +def test_promote_1d_leaves_object_array_roundtrip() -> None: + """``_promote_1d_leaves`` + ``_from_wire`` preserves non-tensor leaves. + + Pins the production TD shape (1D tensor + object array + 2D tensor) + against tensordict 0.12.2 reconstruction bugs that could silently + strip ``NonTensorStack`` / ``NonTensorData`` leaves. Symmetric to + the documented ``.contiguous()`` bug in + ``adapters/transfer_queue.py`` lines 558ā€“562. + """ + from nemo_rl.data_plane.adapters.transfer_queue import ( + _from_wire, + _promote_1d_leaves, + ) + + arr = np.empty(4, dtype=object) + arr[:] = [["a", "b"], ["c"], ["d", "e"], ["f"]] + td = TensorDict( + { + "input_ids": torch.zeros(4, 8, dtype=torch.long), + "input_lengths": torch.tensor([4, 3, 2, 1]), # 1D ā†’ promoted + "content": arr, + }, + batch_size=[4], + ) + promoted = _promote_1d_leaves(td) + assert promoted["input_lengths"].shape == (4, 1) + np.testing.assert_array_equal(promoted["content"], arr) + + restored = _from_wire(promoted) + assert restored["input_lengths"].shape == (4,) + np.testing.assert_array_equal(restored["content"], arr) diff --git a/tests/unit/data_plane/test_writeback_pipeline_e2e.py b/tests/unit/data_plane/test_writeback_pipeline_e2e.py new file mode 100644 index 00000000000..9ee18e121b3 --- /dev/null +++ b/tests/unit/data_plane/test_writeback_pipeline_e2e.py @@ -0,0 +1,129 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Lightweight functional test for the TQ writeback leader gate. + +Pins the contract that ``TQWorkerMixin._write_back`` only fires on the +replica-group leader (``_is_replica_leader`` is True). This is the +``-601 ILLEGAL_CLIENT`` regression boundary: any non-leader sibling that +writes back duplicates the upsert and crashes the mooncake_cpu backend. + +CPU-only, Ray-free ā€” uses :class:`NoOpDataPlaneClient` and a tiny +mixin subclass that fakes ``_is_replica_leader``. +""" + +from __future__ import annotations + +import torch + +from nemo_rl.data_plane import KVBatchMeta +from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient +from nemo_rl.data_plane.worker_mixin import TQWorkerMixin +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + +class _FakeWorker(TQWorkerMixin): + def __init__(self, client: NoOpDataPlaneClient, *, is_leader: bool) -> None: + self._dp_client = client + self._is_leader = is_leader + + def _is_replica_leader(self) -> bool: # type: ignore[override] + return self._is_leader + + +def _seed_partition_with_one_sample(client: NoOpDataPlaneClient) -> KVBatchMeta: + from nemo_rl.data_plane.column_io import write_columns + + client.register_partition( + partition_id="train", + fields=["input_ids", "input_lengths", "prev_logprobs"], + num_samples=1, + consumer_tasks=["train"], + ) + meta = KVBatchMeta( + partition_id="train", + task_name="train", + sample_ids=["s0"], + fields=["input_ids", "input_lengths"], + sequence_lengths=[4], + ) + write_columns( + client, + meta, + { + "input_ids": torch.tensor([[1, 2, 3, 4]], dtype=torch.long), + "input_lengths": torch.tensor([4], dtype=torch.long), + }, + ) + return meta + + +def test_writeback_only_leader_writes(): + """Non-leader sibling write must NOT land ā€” that's the -601 bug class.""" + client = NoOpDataPlaneClient() + meta = _seed_partition_with_one_sample(client) + + leader = _FakeWorker(client, is_leader=True) + sibling = _FakeWorker(client, is_leader=False) + + leader._write_back_result_field( + meta, + BatchedDataDict({"logprobs": torch.zeros(1, 4)}), + result_key="logprobs", + tq_field="prev_logprobs", + ) + sibling._write_back_result_field( + meta, + BatchedDataDict({"logprobs": torch.full((1, 4), 99.0)}), + result_key="logprobs", + tq_field="prev_logprobs", + ) + + fetched = client.get_samples( + sample_ids=meta.sample_ids, + partition_id="train", + select_fields=["prev_logprobs"], + ) + assert torch.allclose(fetched["prev_logprobs"], torch.zeros(1, 4)), ( + "TQ holds a non-leader value ā€” duplicate-writer condition that " + "produces -601 ILLEGAL_CLIENT on the Mooncake backend." + ) + + +def test_writeback_single_worker_default_is_leader(): + """Single-process worker (no TP/CP/PP) is trivially a leader.""" + + class _SingleWorker(TQWorkerMixin): + def __init__(self, client: NoOpDataPlaneClient) -> None: + self._dp_client = client + + def _local_coords(self) -> dict[str, int]: + # No replicated axes ā€” every axis check trivially True. + return {} + + client = NoOpDataPlaneClient() + meta = _seed_partition_with_one_sample(client) + + w = _SingleWorker(client) + w._write_back_result_field( + meta, + BatchedDataDict({"logprobs": torch.full((1, 4), 7.5)}), + result_key="logprobs", + tq_field="prev_logprobs", + ) + fetched = client.get_samples( + sample_ids=meta.sample_ids, + partition_id="train", + select_fields=["prev_logprobs"], + ) + assert torch.allclose(fetched["prev_logprobs"], torch.full((1, 4), 7.5)) diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 797cf9ef914..6ce16df86d6 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -397,3 +397,19 @@ logger: cluster: gpus_per_node: 1 num_nodes: 1 + +# TransferQueue-mediated data plane for sync GRPO. +# Off by default ā€” the legacy grpo_train trainer never engages this. +# Flip enabled=true and run grpo_train_sync to use TQ-mediated bulk +# transfer between rollout and train. See nemo_rl/data_plane/README.md. +data_plane: + enabled: false + impl: transfer_queue + backend: "simple" # TQ storage backend ('simple' or 'mooncake_cpu') + storage_capacity: 1000000 # max samples retained per partition + num_storage_units: 2 # storage shards + claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence + global_segment_size: 549755813888 # 512 GiB ā€” used when backend == "mooncake_cpu" + local_buffer_size: 68719476736 # 64 GiB ā€” used when backend == "mooncake_cpu" + # observability: # NotRequired + # enabled: false diff --git a/uv.lock b/uv.lock index 2fb6cf896b0..e6e30e2296e 100644 --- a/uv.lock +++ b/uv.lock @@ -124,6 +124,7 @@ overrides = [ { name = "flashinfer-python", specifier = ">=0.5.0" }, { name = "llguidance", specifier = ">=1.3.0,<1.4.0" }, { name = "mlflow", specifier = ">=3.11.1" }, + { name = "numpy", specifier = ">=2.1.0" }, { name = "nvidia-cublas", marker = "sys_platform != 'darwin'", specifier = "==13.3.0.5" }, { name = "nvidia-cudnn-cu13", marker = "sys_platform != 'darwin'", specifier = "==9.20.0.48" }, { name = "nvidia-cutlass-dsl", specifier = ">=4.4.1" }, @@ -428,15 +429,18 @@ name = "apache-tvm-ffi" version = "0.1.11" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "platform_machine != 's390x' and sys_platform == 'linux'", - "platform_machine == 's390x' and sys_platform == 'linux'", - "platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'linux'", - "platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'linux'", - "platform_machine != 's390x' and sys_platform == 'darwin'", - "platform_machine == 's390x' and sys_platform == 'darwin'", + "platform_machine != 's390x' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang' and extra != 'extra-7-nemo-rl-vllm'", + "platform_machine == 's390x' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang' and extra != 'extra-7-nemo-rl-vllm'", + "platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang' and extra != 'extra-7-nemo-rl-vllm'", + "platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang' and extra != 'extra-7-nemo-rl-vllm'", + "platform_machine != 's390x' and sys_platform == 'darwin' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang' and extra != 'extra-7-nemo-rl-vllm'", + "platform_machine == 's390x' and sys_platform == 'darwin' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang' and extra != 'extra-7-nemo-rl-vllm'", + "platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang' and extra != 'extra-7-nemo-rl-vllm'", + "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang' and extra != 'extra-7-nemo-rl-vllm'", + "platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang' and extra != 'extra-7-nemo-rl-vllm'", ] dependencies = [ - { name = "typing-extensions", marker = "(extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra != 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "typing-extensions", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-7-nemo-rl-automodel') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6d/3d/4b9226cd45aa800a6904603dda9b323d728f3c3869952a673f3483b78b19/apache_tvm_ffi-0.1.11.tar.gz", hash = "sha256:153cd2c5a9717804cb0bcd9b2709f22a1e5f80ed05b5a490faf5949b136eedba", size = 2798354, upload-time = "2026-05-04T17:48:43.852Z" } wheels = [ @@ -3251,7 +3255,7 @@ name = "ml-dtypes" version = "0.5.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy" }, + { name = "numpy", marker = "(platform_machine != 's390x' and sys_platform == 'linux') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra != 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz", hash = "sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453", size = 692314, upload-time = "2025-11-17T22:32:31.031Z" } wheels = [ @@ -3383,6 +3387,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/cd/07523b9008d5beccebf0fcbcb33b43924bd12dfbbe3b5e4520fdad52aaca/modelscope-1.36.3-py3-none-any.whl", hash = "sha256:65834a077347522d4473778692fded0b23b2a91cb3305811de0deabb83f20e98", size = 6085015, upload-time = "2026-04-28T18:00:50.056Z" }, ] +[[package]] +name = "mooncake-transfer-engine-cuda13" +version = "0.3.10.post2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohttp", marker = "(platform_machine != 'aarch64' and platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 's390x' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "requests", marker = "(platform_machine != 'aarch64' and platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 's390x' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/b9/e6/14538fa71453c4394f8e4dc4a4ba9b90524c011bbaa87db96b5c66ebf09b/mooncake_transfer_engine_cuda13-0.3.10.post2-cp313-cp313-manylinux_2_35_x86_64.whl", hash = "sha256:a96794f4d3c693e6e71ad85ef578a429ec69ab36e0c2f9b45b200d37e45d3cc0", size = 44756026, upload-time = "2026-04-22T03:49:07.836Z" }, + { url = "https://files.pythonhosted.org/packages/87/c9/72353a202de45eceef0525875ac20f3076874e87480c03eb41cc4af37a4e/mooncake_transfer_engine_cuda13-0.3.10.post2-cp313-cp313-manylinux_2_39_aarch64.whl", hash = "sha256:9f70e3aaba4df56fd09e8e4503edc701ac32eedf87641171b4fed344a8ccd0f9", size = 16848965, upload-time = "2026-04-22T06:00:14.941Z" }, +] + [[package]] name = "mpmath" version = "1.3.0" @@ -3813,6 +3830,7 @@ dependencies = [ { name = "math-verify" }, { name = "matplotlib" }, { name = "mlflow" }, + { name = "mooncake-transfer-engine-cuda13", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "nccl4py", marker = "sys_platform != 'darwin' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "ninja" }, { name = "num2words" }, @@ -3834,12 +3852,15 @@ dependencies = [ { name = "swanlab" }, { name = "sympy" }, { name = "tensorboard" }, + { name = "tensordict" }, { name = "tiktoken" }, + { name = "tilelang", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform != 'darwin' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "torchdata" }, { name = "torchvision", version = "0.26.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "torchvision", version = "0.26.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform != 'darwin' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "transferqueue" }, { name = "transformers" }, { name = "triton", version = "3.6.0", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "wandb" }, @@ -3984,6 +4005,7 @@ requires-dist = [ { name = "megatron-bridge", marker = "extra == 'mcore'", editable = "3rdparty/Megatron-Bridge-workspace" }, { name = "megatron-core", marker = "extra == 'mcore'", editable = "3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/3rdparty/Megatron-LM" }, { name = "mlflow", specifier = ">=3.11.1" }, + { name = "mooncake-transfer-engine-cuda13", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')", specifier = "==0.3.10.post2" }, { name = "nccl4py", marker = "sys_platform != 'darwin'" }, { name = "nemo-automodel", extras = ["moe"], marker = "extra == 'automodel'", editable = "3rdparty/Automodel-workspace/Automodel" }, { name = "nemo-gym", marker = "extra == 'nemo-gym'", editable = "3rdparty/Gym-workspace/Gym" }, @@ -4016,12 +4038,15 @@ requires-dist = [ { name = "swanlab" }, { name = "sympy", specifier = ">=1.14.0" }, { name = "tensorboard" }, + { name = "tensordict" }, { name = "tiktoken" }, + { name = "tilelang", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "torch", marker = "sys_platform != 'darwin'", specifier = "==2.11.0", index = "https://download.pytorch.org/whl/cu130" }, { name = "torch", marker = "sys_platform == 'darwin'", specifier = "==2.11.0", index = "https://pypi.org/simple" }, { name = "torchdata" }, { name = "torchvision", marker = "sys_platform != 'darwin'", specifier = "==0.26.0", index = "https://download.pytorch.org/whl/cu130" }, { name = "torchvision", marker = "sys_platform == 'darwin'", specifier = "==0.26.0", index = "https://pypi.org/simple" }, + { name = "transferqueue", git = "https://github.com/Ascend/TransferQueue.git?rev=b266d39" }, { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'automodel'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=v2.14.1" }, { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'mcore'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=v2.14.1" }, { name = "transformers", specifier = "==5.3.0" }, @@ -5748,6 +5773,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/dd/96da98f892250475bdf2328112d7468abdd4acc7b902b6af23f4ed958ea0/pytz-2026.2-py2.py3-none-any.whl", hash = "sha256:04156e608bee23d3792fd45c94ae47fae1036688e75032eea2e3bf0323d1f126", size = 510141, upload-time = "2026-05-04T01:35:27.408Z" }, ] +[[package]] +name = "pyvers" +version = "0.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/32/99/23c73a1298b1c642d8ebdd78e1db4daf1e474152e6839df4f5c93357a3db/pyvers-0.2.2.tar.gz", hash = "sha256:205026bcd0b4c09198cb3a32f243fd179ef012882ce16d93dcb755320acd56f7", size = 12104, upload-time = "2026-01-23T14:12:07.619Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/36/bf/ea840f706b7824dd57220484465995309c8c217995ddb7ce4b262240e912/pyvers-0.2.2-py3-none-any.whl", hash = "sha256:c4696408a0b15fbaa90df33d3bc579cf23a74a73541858f5470216f12f51f3b1", size = 11569, upload-time = "2026-01-23T14:12:06.246Z" }, +] + [[package]] name = "pywin32" version = "311" @@ -6884,6 +6918,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/73/c6/825dab04195756cf8ff2e12698f22513b3db2f64925bdd41671bfb33aaa5/tensorboard_data_server-0.7.2-py3-none-manylinux_2_31_x86_64.whl", hash = "sha256:ef687163c24185ae9754ed5650eb5bc4d84ff257aabdc33f0cc6f74d8ba54530", size = 6590363, upload-time = "2023-10-23T21:23:35.583Z" }, ] +[[package]] +name = "tensordict" +version = "0.12.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cloudpickle" }, + { name = "importlib-metadata" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pyvers" }, + { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform != 'darwin' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/99/e8/ec3f0d5c1c96ff2ffe6eee27030aacf4c863a2d936a7e17fcd1b6cb63c3d/tensordict-0.12.4-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:853b6420c2458861434855453d75052b55887bcca2c4958fe9883813ba30a913", size = 890147, upload-time = "2026-05-22T00:09:29.602Z" }, + { url = "https://files.pythonhosted.org/packages/4f/c3/ae214fbda9f2fe85bca76b272a7924d6a8b58990ba1b167028ae79bc0a85/tensordict-0.12.4-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:3cfd1124b1931780b9e193a9fe7b37d50e5229dae4eaa715db5608c28803a710", size = 533774, upload-time = "2026-05-22T00:09:31.619Z" }, + { url = "https://files.pythonhosted.org/packages/ed/3f/7e7f87da0a343ae234fc346653e812710c0c7823ceb1034b35652f7cbd90/tensordict-0.12.4-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:43e190dc05d217af3d27c207125db90ff5de1a7c5945aab34430a0d5cf81f7fd", size = 537544, upload-time = "2026-05-22T00:09:33.524Z" }, + { url = "https://files.pythonhosted.org/packages/77/0a/b765ae434ef1650b3f538fdc5ec979b2188a2c3e839a6dddb3b173f6d033/tensordict-0.12.4-cp313-cp313-win_amd64.whl", hash = "sha256:0d96da5907b7a5dbd10782a4166eb0e82a702e805b11f94a28bd629da61dff36", size = 586791, upload-time = "2026-05-22T00:09:35.659Z" }, + { url = "https://files.pythonhosted.org/packages/b3/84/c84936bdc4c2d1432f96d4e16f2521e196208332f985de6329bb8398d127/tensordict-0.12.4-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:a1e23296684e532650e236228c59fe0f4dd323d7c409c0798c18fd2791c1e252", size = 895573, upload-time = "2026-05-22T00:09:37.429Z" }, + { url = "https://files.pythonhosted.org/packages/6d/b6/d574e2b758631563861d51cba4cc595d27a3965db3473a05ab268eead05b/tensordict-0.12.4-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:6e60888bc24990ead02d52f16fa607af8c01c92089ad767540eca88ade5fb49f", size = 535213, upload-time = "2026-05-22T00:09:39.561Z" }, + { url = "https://files.pythonhosted.org/packages/13/a4/25c29e653878e58ed3cb111146e4dd8cdb4cfd4b6f66dd2080f94f8e78f4/tensordict-0.12.4-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:031c70d2101376e0fb8036b017c8271a27892c1b9ba6aea021c039c7535aac53", size = 539088, upload-time = "2026-05-22T00:09:41.377Z" }, + { url = "https://files.pythonhosted.org/packages/35/6f/c8107ea679a60e7584bc6d36b854879a33f5a990819e174a7ed653edb781/tensordict-0.12.4-cp313-cp313t-win_amd64.whl", hash = "sha256:a1320ea2ed9e0289209b0efc51b8bf2bca02cf5273fade3aec4f60a4ddfed61b", size = 597644, upload-time = "2026-05-22T00:09:42.922Z" }, +] + [[package]] name = "threadpoolctl" version = "3.6.0" @@ -6924,18 +6982,19 @@ name = "tilelang" version = "0.1.9" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "apache-tvm-ffi", version = "0.1.9", source = { registry = "https://pypi.org/simple" } }, - { name = "cloudpickle" }, - { name = "ml-dtypes" }, - { name = "numpy" }, - { name = "psutil" }, - { name = "setuptools" }, + { name = "apache-tvm-ffi", version = "0.1.9", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "apache-tvm-ffi", version = "0.1.11", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-7-nemo-rl-automodel') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "cloudpickle", marker = "(platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 's390x' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "ml-dtypes", marker = "(platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 's390x' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "numpy", marker = "(platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 's390x' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "psutil", marker = "(platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 's390x' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "setuptools", marker = "(platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 's390x' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (sys_platform == 'darwin' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, - { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "(sys_platform != 'darwin' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (sys_platform != 'darwin' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, - { name = "torch-c-dlpack-ext" }, - { name = "tqdm" }, - { name = "typing-extensions" }, - { name = "z3-solver" }, + { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "(platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 's390x' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra != 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang')" }, + { name = "torch-c-dlpack-ext", marker = "(platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 's390x' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "tqdm", marker = "(platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 's390x' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "typing-extensions", marker = "(platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 's390x' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "z3-solver", marker = "(platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 's390x' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 's390x' and sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 's390x' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra != 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/56/70/5051f65821baa30a3d61fc48f8ba10c776490315e8c90f82559b92089756/tilelang-0.1.9.tar.gz", hash = "sha256:287f727c913bb648fcf6c1968809ba3390e55eeed257a5c6bb9a80bc05966af4", size = 93395292, upload-time = "2026-04-22T09:19:11.988Z" } wheels = [ @@ -7118,7 +7177,7 @@ version = "0.1.5" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'darwin' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm')" }, - { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "(sys_platform != 'darwin' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "torch", version = "2.11.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform == 'darwin' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore') or (sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform == 'linux' and extra != 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/37/de/921b6491efce5c389a5ef9bbed3d2d6660005840dae488124173180859ab/torch_c_dlpack_ext-0.1.5.tar.gz", hash = "sha256:d06f0357d575d22a168cc77acb9020fc4bae30968ceb6718a055dcbe92bacabe", size = 12913, upload-time = "2026-01-12T11:25:08.484Z" } wheels = [ @@ -7361,6 +7420,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/da/98/a9937a969d018a23badfea0b381f66783649d48e0ea6c41923265c3cbeb3/traitlets-5.15.0-py3-none-any.whl", hash = "sha256:fb36a18867a6803deab09f3c5e0fa81bb7b26a5c9e82501c9933f759166eff40", size = 85877, upload-time = "2026-05-06T08:05:55.853Z" }, ] +[[package]] +name = "transferqueue" +version = "0.1.7.dev0" +source = { git = "https://github.com/Ascend/TransferQueue.git?rev=b266d39#b266d39a15aae114730de36cf8317b6285436f7f" } +dependencies = [ + { name = "hydra-core" }, + { name = "msgspec" }, + { name = "numpy" }, + { name = "omegaconf" }, + { name = "psutil" }, + { name = "pyzmq" }, + { name = "ray", extra = ["default"] }, + { name = "tensordict" }, +] + [[package]] name = "transformer-engine" version = "2.14.1+366798e"