Skip to content

perf(checkpoint): batch generation-prefix TQ writes #4389

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
macandro96 wants to merge 2 commits into from
Closed

perf(checkpoint): batch generation-prefix TQ writes #4389

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2c646aa
perf(checkpoint): batch worker prefix writes and add scale benchmarks
Oct 1, 2026
9812a07
refactor(checkpoint): share prefix batch limits
macandro96 Oct 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions examples/configs/grpo_math_1B_megatron_single_controller.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,8 @@ rollout_checkpointing:
# reaches this many tokens, rather than only when a checkpoint cut fires.
# Bounds crash exposure independently of checkpoint cadence; 0 disables it.
generation_chunk_flush_tokens: 0
generation_prefix_batch_size: 256 # Checkpoint PUT rows per generation owner; 1 disables batching.
generation_prefix_batch_max_tokens: 4194304 # Aggregate tokens; one oversized row goes alone.
prepare_timeout_s: 300

policy:
Expand Down
2 changes: 2 additions & 0 deletions examples/configs/ppo_math_1B_megatron_single_controller.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,8 @@ rollout_checkpointing:
# reaches this many tokens, rather than only when a checkpoint cut fires.
# Bounds crash exposure independently of checkpoint cadence; 0 disables it.
generation_chunk_flush_tokens: 0
generation_prefix_batch_size: 256 # Checkpoint PUT rows per generation owner; 1 disables batching.
generation_prefix_batch_max_tokens: 4194304 # Aggregate tokens; one oversized row goes alone.
prepare_timeout_s: 300

policy:
Expand Down
7 changes: 7 additions & 0 deletions nemo_rl/algorithms/single_controller_utils/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -759,6 +759,13 @@ class GymRolloutCheckpointConfig(BaseModel, extra="forbid"):
# Stage an in-flight call's accumulated tokens once its unstaged segment
# reaches this many tokens. 0 disables periodic chunk flushing.
generation_chunk_flush_tokens: Annotated[int, Field(ge=0)] = 0
# Checkpoint-triggered prefix PUTs per generation owner. 1 retains the
# single-row path. Periodic chunk flushing is unchanged.
generation_prefix_batch_size: Annotated[int, Field(gt=0)] = 256
# Aggregate staged tokens per batch (~64 MiB of base tensors at default).
# Not a total-memory bound: metadata/copies add overhead; one oversized
# record is allowed alone, since a record cannot be split here.
generation_prefix_batch_max_tokens: Annotated[int, Field(gt=0)] = 4_194_304
prepare_timeout_s: Annotated[float, Field(gt=0)] = 300.0

@property
Expand Down
2 changes: 2 additions & 0 deletions nemo_rl/algorithms/single_controller_utils/setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2172,6 +2172,8 @@ def _build_trainer_then_megatron_generation() -> tuple[
rollout_checkpoint_cfg.gym.generation_prefix_cuts_enabled
),
generation_cut_control_token=token_capture_cfg.control_auth_token,
generation_prefix_batch_size=rollout_checkpoint_cfg.gym.generation_prefix_batch_size,
generation_prefix_batch_max_tokens=rollout_checkpoint_cfg.gym.generation_prefix_batch_max_tokens,
generation_chunk_flush_tokens=(
rollout_checkpoint_cfg.gym.generation_chunk_flush_tokens
),
Expand Down
323 changes: 205 additions & 118 deletions nemo_rl/data_plane/tq_token_sink.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@

import json
import logging
from collections.abc import Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

Expand All @@ -53,6 +54,7 @@
StageResult,
)

from nemo_rl.data_plane.codec import stack_or_nest
from nemo_rl.data_plane.schema import (
ROUTE_ENCODING_ENVELOPE,
ROUTE_ENCODING_LIST,
Expand Down Expand Up @@ -240,124 +242,7 @@ def _stage_at_key(self, record: StagedCallRecord, key: str) -> StageResult:
from nemo_gym.token_id_capture.staging.records import StageResult

try:
field_dict = {
"token_ids_delta": torch.tensor(
[record.token_ids_delta], dtype=torch.int64
),
"token_mask_delta": torch.tensor(
[record.token_mask_delta], dtype=torch.float32
),
"generation_logprobs_delta": torch.tensor(
[record.generation_log_probs_delta], dtype=torch.float32
),
"schema_version": torch.tensor(
[record.schema_version], dtype=torch.int64
),
"digest_version": torch.tensor(
[record.digest_version], dtype=torch.int64
),
"extras_digest_version": torch.tensor(
[record.extras_digest_version], dtype=torch.int64
),
"rollout_id_utf8": _bytes_tensor(record.rollout_id.encode("utf-8")),
"model_call_id_utf8": _bytes_tensor(
record.model_call_id.encode("utf-8")
),
"parent_call_id_utf8": _bytes_tensor(
(record.parent_call_id or "\0").encode("utf-8")
),
"parent_call_id_present": torch.tensor(
[record.parent_call_id is not None], dtype=torch.bool
),
"capture_mode": torch.tensor(
[_MODE_TO_CODE[record.mode]], dtype=torch.int64
),
"prev_len": torch.tensor([record.prev_len], dtype=torch.int64),
"delta_len": torch.tensor([record.delta_len], dtype=torch.int64),
"cum_len": torch.tensor([record.cum_len], dtype=torch.int64),
"weight_version": torch.tensor(
[record.weight_version], dtype=torch.int64
),
"digest_bytes": _bytes_tensor(bytes.fromhex(record.digest)),
"extras_digest_bytes": _bytes_tensor(
bytes.fromhex(record.extras_digest)
),
}
chain_hash, chain_hash_present = _optional_digest_fields(record.chain_hash)
cumulative_hash, cumulative_hash_present = _optional_digest_fields(
record.cumulative_hash
)
field_dict.update(
{
"chain_hash_bytes": chain_hash,
"chain_hash_present": chain_hash_present,
"cumulative_hash_bytes": cumulative_hash,
"cumulative_hash_present": cumulative_hash_present,
}
)
extras_metadata = dict(record.extras) if record.extras is not None else None
routed = (
extras_metadata.pop("routed_experts", None)
if extras_metadata is not None
else None
)
field_dict[ROUTED_EXTRAS_METADATA_FIELD] = _bytes_tensor(
json.dumps(
extras_metadata,
sort_keys=True,
separators=(",", ":"),
ensure_ascii=False,
).encode("utf-8")
)
routed_len = 0
routed_encoding = ROUTE_ENCODING_NONE
if routed is not None:
delta_len = len(record.token_ids_delta)
if isinstance(routed, str):
from nemo_rl.utils.routed_experts_codec import (
decode_routed_experts,
)

dtype_name = routed.split(":", 3)[1]
dtype = {
"int8": torch.int8,
"int16": torch.int16,
"int32": torch.int32,
}.get(dtype_name)
if dtype is None:
raise ValueError(
f"unsupported routed_experts dtype {dtype_name!r}"
)
experts = decode_routed_experts(routed, dtype)
routed_encoding = ROUTE_ENCODING_ENVELOPE
else:
experts = torch.tensor(routed, dtype=torch.int16)
routed_encoding = ROUTE_ENCODING_LIST
if experts.dim() != 3 or experts.shape[0] != delta_len:
raise ValueError(
"routed_experts must already be delta-aligned: "
f"got shape {tuple(experts.shape)} for delta_len={delta_len}"
)
field_dict[ROUTED_EXPERTS_FIELD] = experts.unsqueeze(0)
routed_len = int(experts.shape[0])
field_dict[ROUTED_EXPERTS_ENCODING_FIELD] = torch.tensor(
[routed_encoding], dtype=torch.int64
)
field_dict[ROUTED_LEN_FIELD] = torch.tensor([routed_len], dtype=torch.int64)
fields = TensorDict(field_dict, batch_size=[1])
tags = [
{
"rollout_id": record.rollout_id,
"model_call_id": record.model_call_id,
"parent_call_id": record.parent_call_id,
"prev_len": record.prev_len,
"delta_len": record.delta_len,
"cum_len": record.cum_len,
"weight_version": record.weight_version,
"digest": record.digest,
"schema_version": record.schema_version,
}
]
fields, tags = self._encode_record(record)
_call_dp(
self._dp_client,
"put_samples",
Expand All @@ -381,6 +266,208 @@ def _stage_at_key(self, record: StagedCallRecord, key: str) -> StageResult:
)
return StageResult(ok=True, staging_key=key)

@staticmethod
def _encode_record(
record: StagedCallRecord,
) -> tuple[TensorDict, list[dict[str, Any]]]:
"""Encode one row without publishing it."""
field_dict = {
"token_ids_delta": torch.tensor(
[record.token_ids_delta], dtype=torch.int64
),
"token_mask_delta": torch.tensor(
[record.token_mask_delta], dtype=torch.float32
),
"generation_logprobs_delta": torch.tensor(
[record.generation_log_probs_delta], dtype=torch.float32
),
"schema_version": torch.tensor([record.schema_version], dtype=torch.int64),
"digest_version": torch.tensor([record.digest_version], dtype=torch.int64),
"extras_digest_version": torch.tensor(
[record.extras_digest_version], dtype=torch.int64
),
"rollout_id_utf8": _bytes_tensor(record.rollout_id.encode("utf-8")),
"model_call_id_utf8": _bytes_tensor(record.model_call_id.encode("utf-8")),
"parent_call_id_utf8": _bytes_tensor(
(record.parent_call_id or "\0").encode("utf-8")
),
"parent_call_id_present": torch.tensor(
[record.parent_call_id is not None], dtype=torch.bool
),
"capture_mode": torch.tensor(
[_MODE_TO_CODE[record.mode]], dtype=torch.int64
),
"prev_len": torch.tensor([record.prev_len], dtype=torch.int64),
"delta_len": torch.tensor([record.delta_len], dtype=torch.int64),
"cum_len": torch.tensor([record.cum_len], dtype=torch.int64),
"weight_version": torch.tensor([record.weight_version], dtype=torch.int64),
"digest_bytes": _bytes_tensor(bytes.fromhex(record.digest)),
"extras_digest_bytes": _bytes_tensor(bytes.fromhex(record.extras_digest)),
}
chain_hash, chain_hash_present = _optional_digest_fields(record.chain_hash)
cumulative_hash, cumulative_hash_present = _optional_digest_fields(
record.cumulative_hash
)
field_dict.update(
{
"chain_hash_bytes": chain_hash,
"chain_hash_present": chain_hash_present,
"cumulative_hash_bytes": cumulative_hash,
"cumulative_hash_present": cumulative_hash_present,
}
)
extras_metadata = dict(record.extras) if record.extras is not None else None
routed = (
extras_metadata.pop("routed_experts", None)
if extras_metadata is not None
else None
)
field_dict[ROUTED_EXTRAS_METADATA_FIELD] = _bytes_tensor(
json.dumps(
extras_metadata,
sort_keys=True,
separators=(",", ":"),
ensure_ascii=False,
).encode("utf-8")
)
routed_len = 0
routed_encoding = ROUTE_ENCODING_NONE
if routed is not None:
delta_len = len(record.token_ids_delta)
if isinstance(routed, str):
from nemo_rl.utils.routed_experts_codec import (
decode_routed_experts,
)

dtype_name = routed.split(":", 3)[1]
dtype = {
"int8": torch.int8,
"int16": torch.int16,
"int32": torch.int32,
}.get(dtype_name)
if dtype is None:
raise ValueError(f"unsupported routed_experts dtype {dtype_name!r}")
experts = decode_routed_experts(routed, dtype)
routed_encoding = ROUTE_ENCODING_ENVELOPE
else:
experts = torch.tensor(routed, dtype=torch.int16)
routed_encoding = ROUTE_ENCODING_LIST
if experts.dim() != 3 or experts.shape[0] != delta_len:
raise ValueError(
"routed_experts must already be delta-aligned: "
f"got shape {tuple(experts.shape)} for delta_len={delta_len}"
)
field_dict[ROUTED_EXPERTS_FIELD] = experts.unsqueeze(0)
routed_len = int(experts.shape[0])
field_dict[ROUTED_EXPERTS_ENCODING_FIELD] = torch.tensor(
[routed_encoding], dtype=torch.int64
)
field_dict[ROUTED_LEN_FIELD] = torch.tensor([routed_len], dtype=torch.int64)
fields = TensorDict(field_dict, batch_size=[1])
tags = [
{
"rollout_id": record.rollout_id,
"model_call_id": record.model_call_id,
"parent_call_id": record.parent_call_id,
"prev_len": record.prev_len,
"delta_len": record.delta_len,
"cum_len": record.cum_len,
"weight_version": record.weight_version,
"digest": record.digest,
"schema_version": record.schema_version,
}
]
return fields, tags

def stage_generation_prefix_batch(
self,
records: Sequence[StagedCallRecord],
*,
checkpoint_id: str,
chunk_sequences: Sequence[int],
) -> list[StageResult]:
"""Synchronously publish prefixes, preserving one result per input row.

Success is returned only after TQ acknowledges the whole compatible
batch. A failed write reports failure for every member of that batch;
callers must not publish its lineage. Optional route columns and route
shapes are grouped separately so missing routes never become fabricated
payloads. Variable token and identity lengths remain jagged.
"""
# Gym is an optional dependency outside capture-enabled runs.
from nemo_gym.token_id_capture.staging.records import StageResult

if len(records) != len(chunk_sequences):
raise ValueError("records and chunk_sequences must have equal lengths")
keys = [
generation_cut_staging_key(
checkpoint_id,
record.rollout_id,
record.model_call_id,
chunk_sequence=sequence,
)
for record, sequence in zip(records, chunk_sequences, strict=True)
]
if len(set(keys)) != len(keys):
raise ValueError("generation-prefix batch contains duplicate staging keys")
if not records:
return []

results = [
StageResult(ok=False, staging_key=key, error="not written") for key in keys
]
groups: dict[tuple[Any, ...], list[tuple[int, TensorDict, dict[str, Any]]]] = {}
for index, record in enumerate(records):
try:
fields, tags = self._encode_record(record)
signature = tuple(
(name, fields[name].dtype, tuple(fields[name].shape[2:]))
for name in sorted(fields.keys())
)
groups.setdefault(signature, []).append((index, fields, tags[0]))
except Exception as error: # Encoding failures must fail closed.
results[index] = StageResult(
ok=False,
staging_key=keys[index],
error=f"{type(error).__name__}: {error}",
)

for group in groups.values():
indices = [item[0] for item in group]
try:
fields = TensorDict(
{
name: stack_or_nest([item[1][name][0] for item in group])
for name in group[0][1].keys()
},
batch_size=[len(group)],
)
_call_dp(
self._dp_client,
"put_samples",
sample_ids=[keys[index] for index in indices],
partition_id=self._staging_partition,
fields=fields,
tags=[item[2] for item in group],
)
except Exception as error: # Batch failures must fail closed.
logging.getLogger(__name__).warning(
"TQ generation-prefix batch failed for %d rows: %s: %s",
len(group),
type(error).__name__,
error,
)
for index in indices:
results[index] = StageResult(
ok=False,
staging_key=keys[index],
error=f"{type(error).__name__}: {error}",
)
else:
for index in indices:
results[index] = StageResult(ok=True, staging_key=keys[index])
return results

def clear(self, staging_keys: list[str]) -> None:
"""Drop staged rows (finalizer / eviction cleanup)."""
if not staging_keys:
Expand Down
Loading
Loading