Skip to content

Assemble packed training chunks in numpy instead of per-bin jnp ops #2087

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

Open
gutianyu-google wants to merge 1 commit into
base: main
Choose a base branch
from
+83 −134
Open

Assemble packed training chunks in numpy instead of per-bin jnp ops #2087

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
217 changes: 83 additions & 134 deletions tunix/rl/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -513,159 +513,108 @@ def pack_sequences(
# Real segments per row are bounded by the token budget (each segment >= 1
# token). `None` uses that safe bound so `num_segments = budget + 1` never
# overflows; a smaller override shrinks the loss buckets and is enforced by
# the raise in `_flush_pack`.
# the raise in `_emit`.
effective_max_segments = (
max_segments_per_packed_row
if max_segments_per_packed_row is not None
else max_token_budget
)

def _flush_pack(pack_items, example_cls, first_item) -> common.TrainExample:
first_item = first_item or {}
def _emit(chunk):
"""Merges one chunk (pack_size bins) into a [pack_size, budget] example.

Every row is assembled in numpy and moved to device once per feature.
Building each bin with its own jnp ops and concatenating afterwards costs
~18 device dispatches per bin, which with hundreds of bins per step made
the packer itself a visible part of the training step.
"""
first_item = first_item_for_dummy or {}
has_policy_version = first_item.get("policy_version") is not None
kwargs = {}
tracked_per_token_keys = []

if first_item.get("ref_per_token_logps") is not None:
tracked_per_token_keys.append("ref_per_token_logps")
if first_item.get("old_per_token_logps") is not None:
tracked_per_token_keys.append("old_per_token_logps")
if first_item.get("returns") is not None:
tracked_per_token_keys.append("returns")
if first_item.get("old_values") is not None:
tracked_per_token_keys.append("old_values")

if not pack_items:
p_ids_arr = jnp.zeros((1, 0), dtype=jnp.int32)
p_mask_arr = jnp.zeros((1, 0), dtype=jnp.int32)
c_ids_arr = jnp.full((1, max_token_budget), pad_id, dtype=jnp.int32)
c_mask_arr = jnp.zeros((1, max_token_budget), dtype=jnp.int32)
adv_arr = jnp.zeros((1, max_token_budget), dtype=jnp.float32)
seg_arr = jnp.zeros((1, max_token_budget), dtype=jnp.int32)
pos_arr = jnp.zeros((1, max_token_budget), dtype=jnp.int32)

kwargs.update(
prompt_ids=p_ids_arr,
prompt_mask=p_mask_arr,
completion_ids=c_ids_arr,
completion_mask=c_mask_arr,
advantages=adv_arr,
ref_per_token_logps=None,
old_per_token_logps=None,
segment_ids=seg_arr,
segment_positions=pos_arr,
)
for k in tracked_per_token_keys:
kwargs[k] = jnp.zeros((1, max_token_budget), dtype=jnp.float32)
if has_policy_version:
kwargs["policy_version"] = first_item["policy_version"]
return example_cls(**kwargs) # pytype: disable=wrong-keyword-args

# `len(pack_items)` is the real segment count of this row. It cannot exceed
# the token budget (each segment >= 1 token), so the default bound never
# trips; a too-small `max_segments_per_packed_row` override does, and we
# fail loud rather than let `segment_sum` silently drop the overflow.
if len(pack_items) > effective_max_segments:
raise ValueError(
f"pack_sequences: a packed row has {len(pack_items)} segments, "
f"exceeding max_segments_per_packed_row={effective_max_segments}; "
"increase it (or leave it None for the budget-derived safe default)."
)
tracked_per_token_keys = [
k
for k in (
"ref_per_token_logps",
"old_per_token_logps",
"returns",
"old_values",
)
if first_item.get(k) is not None
]

current_tokens = sum(
len(it["prompt_ids"]) + len(it["completion_ids"]) for it in pack_items
)
pad_len = max_token_budget - current_tokens

packed_c_ids = []
packed_c_mask = []
packed_adv = []
packed_segment_ids = []
packed_positions = []

per_token_feature_buffers = {k: [] for k in tracked_per_token_keys}

for i, item in enumerate(pack_items, start=1):
p_ids = item["prompt_ids"]
c_ids = item["completion_ids"]
seq_len = len(p_ids) + len(c_ids)

packed_c_ids.extend([p_ids, c_ids])
packed_c_mask.extend([np.zeros_like(p_ids), item["completion_mask"]])

if item["adv_is_per_token"]:
packed_adv.extend([
np.zeros_like(p_ids, dtype=np.float32),
item["advantages"],
])
else:
packed_adv.extend([
np.zeros_like(p_ids, dtype=np.float32),
np.full(len(c_ids), item["advantages"], dtype=np.float32),
])

for k in tracked_per_token_keys:
per_token_feature_buffers[k].extend([
np.zeros_like(p_ids, dtype=np.float32),
item[k],
])

packed_segment_ids.append(np.full(seq_len, i, dtype=np.int32))
packed_positions.append(np.arange(seq_len, dtype=np.int32))

def _pad(arr_list, val, length):
arr = np.concatenate(arr_list) if arr_list else np.array([])
return np.pad(arr, (0, length), constant_values=val)

p_ids_arr = jnp.zeros((1, 0), dtype=jnp.int32)
p_mask_arr = jnp.zeros((1, 0), dtype=jnp.int32)

c_ids_arr = jnp.array(_pad(packed_c_ids, pad_id, pad_len))[None, :]
c_mask_arr = jnp.array(_pad(packed_c_mask, 0, pad_len))[None, :]
adv_arr = jnp.array(_pad(packed_adv, 0.0, pad_len))[None, :]
seg_arr = jnp.array(_pad(packed_segment_ids, 0, pad_len))[None, :]
pos_arr = jnp.array(_pad(packed_positions, 0, pad_len))[None, :]

per_token_features = {}
for k in tracked_per_token_keys:
per_token_features[k] = jnp.array(
_pad(per_token_feature_buffers[k], 0.0, pad_len)
)[None, :]
num_rows = len(chunk)
budget = max_token_budget
c_ids = np.full((num_rows, budget), pad_id, dtype=np.int32)
c_mask = np.zeros((num_rows, budget), dtype=np.int32)
adv = np.zeros((num_rows, budget), dtype=np.float32)
seg = np.zeros((num_rows, budget), dtype=np.int32)
pos = np.zeros((num_rows, budget), dtype=np.int32)
per_token_features = {
k: np.zeros((num_rows, budget), dtype=np.float32)
for k in tracked_per_token_keys
}

for row, pack_items in enumerate(chunk):
# `len(pack_items)` is the real segment count of this row. It cannot
# exceed the token budget (each segment >= 1 token), so the default
# bound never trips; a too-small `max_segments_per_packed_row` override
# does, and we fail loud rather than let `segment_sum` silently drop the
# overflow. An empty bin stays a dummy row of padding.
if len(pack_items) > effective_max_segments:
raise ValueError(
f"pack_sequences: a packed row has {len(pack_items)} segments, "
f"exceeding max_segments_per_packed_row={effective_max_segments}; "
"increase it (or leave it None for the budget-derived safe"
" default)."
)
offset = 0
for i, item in enumerate(pack_items, start=1):
p_ids = item["prompt_ids"]
c_ids_item = item["completion_ids"]
prompt_len = len(p_ids)
seq_len = prompt_len + len(c_ids_item)
completion = slice(offset + prompt_len, offset + seq_len)
c_ids[row, offset : offset + prompt_len] = p_ids
c_ids[row, completion] = c_ids_item
c_mask[row, completion] = item["completion_mask"]
# A scalar advantage broadcasts over the completion tokens.
adv[row, completion] = item["advantages"]
for k in tracked_per_token_keys:
per_token_features[k][row, completion] = item[k]
seg[row, offset : offset + seq_len] = i
pos[row, offset : offset + seq_len] = np.arange(seq_len, dtype=np.int32)
offset += seq_len

kwargs = dict(
prompt_ids=p_ids_arr,
prompt_mask=p_mask_arr,
completion_ids=c_ids_arr,
completion_mask=c_mask_arr,
advantages=adv_arr,
prompt_ids=jnp.zeros((num_rows, 0), dtype=jnp.int32),
prompt_mask=jnp.zeros((num_rows, 0), dtype=jnp.int32),
completion_ids=jnp.asarray(c_ids),
completion_mask=jnp.asarray(c_mask),
advantages=jnp.asarray(adv),
ref_per_token_logps=None,
old_per_token_logps=None,
segment_ids=seg_arr,
segment_positions=pos_arr,
segment_ids=jnp.asarray(seg),
segment_positions=jnp.asarray(pos),
)
for k in tracked_per_token_keys:
kwargs[k] = per_token_features[k]

kwargs[k] = jnp.asarray(per_token_features[k])
if has_policy_version:
kwargs["policy_version"] = pack_items[0]["policy_version"]

# Each row carries the policy version of its first sequence; a dummy
# (empty) row inherits the first item's.
kwargs["policy_version"] = jnp.concatenate(
[
jnp.asarray(
pack_items[0]["policy_version"]
if pack_items
else first_item["policy_version"]
)
for pack_items in chunk
],
axis=0,
)
return example_cls(**kwargs) # pytype: disable=wrong-keyword-args

chunk_capacity = pack_size * max_token_budget

def _emit(chunk):
"""Merges one chunk (pack_size bins) into a [pack_size, budget] example."""
chunk_examples = [
_flush_pack(bin_items, example_cls, first_item_for_dummy)
for bin_items in chunk
]
return jax.tree.map(
lambda first_x, *rest_xs: None
if first_x is None
else jnp.concatenate((first_x, *rest_xs), axis=0),
*chunk_examples,
)

def _mark(merged, is_update):
# `num_segments = effective_max_segments + 1` (+1 = padding bucket) is a
# static upper bound, fixed every step so the segment-aware loss compiles
Expand Down