Skip to content

trainer_rank: retain completed split-forward memory peaks #894

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

Merged
bradhilton merged 6 commits into from
Sep 15, 2026
+802 −58
Merged

trainer_rank: retain completed split-forward memory peaks #894

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
f879d4b
trainer_rank: price enclosing routed-expert buffers at FC2
bradhilton Sep 14, 2026
b2e4050
fix(trainer-rank): retain observed memory estimates beyond calibratio…
bradhilton Sep 14, 2026
fe97538
Remove obsolete post-calibration cost-drop search
bradhilton Sep 14, 2026
0771116
Retain completed split-forward memory peaks for admission
bradhilton Sep 14, 2026
4c8cb6a
Cover incomplete split-caller memory attribution
bradhilton Sep 14, 2026
4d9d788
Fix split-memory test fixture type annotations
bradhilton Sep 14, 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
261 changes: 217 additions & 44 deletions src/art/trainer_rank/_impl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1205,7 +1205,7 @@ def _configure_moe_dispatcher_caches(model: Sequence[torch.nn.Module]) -> None:
def _moe_output_bytes_per_token(
model: Sequence[torch.nn.Module], shape: ParallelShape
) -> int:
"""Known eager FC2 input/outputs, not the complete or compiled working set."""
"""Known routed-expert working set, not a complete model/compiled bound."""
if shape != ParallelShape(tp=1, cp=1):
return 0
from megatron.core.extensions.transformer_engine import TERowParallelGroupedLinear
Expand All @@ -1216,7 +1216,7 @@ def _moe_output_bytes_per_token(
MoEAlltoAllTokenDispatcher,
)

from art.megatron.lora import LoRA, MLPExpertsLinearFC2LoRA
from art.megatron.lora import LoRA, MLPExpertsLinearFC1LoRA, MLPExpertsLinearFC2LoRA

coefficient = 0
for chunk in model:
Expand Down Expand Up @@ -1288,6 +1288,27 @@ def _moe_output_bytes_per_token(
# their sum. Compilers may reuse storage; other workspace is
# not covered. Unknown input metadata keeps the prior pair.
features = inputs.shape[-2] + 3 * fc2.out_features
fc1 = getattr(experts, "linear_fc1", None)
if (
type(fc1) is MLPExpertsLinearFC1LoRA
and "forward" not in vars(fc1)
and not getattr(fc1, "_forward_hooks", None)
and not getattr(fc1, "_forward_pre_hooks", None)
and fc1.fused_gate_up
and not fc1.non_gated
and fc1.out_features == 2 * inputs.shape[-2]
and getattr(dispatcher, "ep_size", None) == 1
and getattr(dispatcher, "tp_size", None) == 1
and getattr(dispatcher, "num_local_experts", 0) > 1
and getattr(config, "moe_permute_fusion", False)
and getattr(experts, "offload_expert_fc1", None) is False
and getattr(experts, "offload_moe_act", None) is False
and getattr(experts, "activation_recompute", None) is False
):
# The two dispatched H-wide inputs and FC1 gate/up sum
# remain live at the FC2 sum, including in the observed
# compiled path. This is one stage, not a backward bound.
features += 2 * fc2.out_features + fc1.out_features
coefficient = max(
coefficient,
config.moe_router_topk * features * weights.element_size(),
Expand Down Expand Up @@ -1422,6 +1443,8 @@ def __init__(self, runtime: TrainingRuntime) -> None:
self._hybridep_buffer_id: int | None = None
self._hybridep_rows_high_water = 0
self._memory_profiles: dict[_MemorySignature, _MemoryProfile] = {}
self._split_memory_floors: dict[bytes, int] = {}
self._split_memory_floor_status = "not_observed"
self._last_global_micro_batch_size: int | None = None
self._skipped_forward_waves: dict[object, tuple[int, int, int]] = {}
# Bounded LRU: steady-state hits are temporally local (identical
Expand Down Expand Up @@ -2271,12 +2294,13 @@ def _forward_micro_batches(
)
)
else:
tracked_outputs = self._execute_admitted_plan(
candidate.plan,
check=candidate.check,
context="forward_micro_batches",
tracked_outputs, memory_baseline, forward_peak = (
self._execute_split_plan_with_memory_tracking(
candidate.plan,
check=candidate.check,
context="forward_micro_batches",
)
)
memory_baseline = None
flat_outputs = iter(tracked_outputs)
outputs = [_unflatten(item, flat_outputs) for item in candidate.inputs]
stop = start + candidate.stats_global_count
Expand Down Expand Up @@ -2324,6 +2348,10 @@ def _forward_micro_batches(
# to the forward's return, already recorded for this same plan.
if isinstance(candidate.plan, _FlatForwardPlan):
self._update_peak_memory_profile(candidate.plan, memory_baseline)
elif memory_baseline is not None:
self._record_split_memory_floor(
candidate.plan, memory_baseline, forward_peak
)
# Only the caller may retain completed outputs into the next wave.
del tracked_outputs, flat_outputs, outputs
start = stop
Expand Down Expand Up @@ -2411,14 +2439,27 @@ def _execute_admitted_plan(
plan, check=check, context=context
)
return outputs
outputs, _baseline, _peak = self._execute_split_plan_with_memory_tracking(
plan, check=check, context=context
)
return outputs

def _execute_split_plan_with_memory_tracking(
self, plan: _SplitForwardPlan, *, check: _MemoryCheck, context: str
) -> tuple[list[AnyForwardOutput], int | None, int]:
baseline, peak = None, 0
merged: list[AnyForwardOutput | None] = [None] * plan.request_count
for ordinal, (subforward, indices) in enumerate(
zip(plan.subforwards, plan.request_indices, strict=True)
):
try:
outputs, _baseline = self._run_flat_plan_with_memory_tracking(
outputs, child_baseline = self._run_flat_plan_with_memory_tracking(
subforward, check=check, context=context
)
if child_baseline is not None:
if baseline is None:
baseline = child_baseline
peak = max(peak, int(torch.cuda.max_memory_allocated(self.device)))
except TrainerRankMemoryError as error:
# Model execution already began, so no replanning is possible
# and the caller must not mistake this for an up-front refusal.
Expand All @@ -2434,7 +2475,7 @@ def _execute_admitted_plan(
merged[index] = output
if any(output is None for output in merged):
raise AssertionError("split execution did not cover every request")
return cast(list[AnyForwardOutput], merged)
return cast(list[AnyForwardOutput], merged), baseline, peak

def _plan_admissible_forward(
self,
Expand Down Expand Up @@ -2589,28 +2630,177 @@ def _admit_split_rung(
for chunk in chunks
]
costs = [self._plan_cost(plan) for plan in plans]
check = self._split_rung_check(costs)
# Bind original request mappings; floor keys normalize execution order.
order = sorted(range(len(plans)), key=lambda i: (-costs[i].ephemeral, i))
split = _SplitForwardPlan(
subforwards=tuple(plans[i] for i in order),
request_indices=tuple(tuple(chunks[i]) for i in order),
request_count=len(requests),
)
check = self._split_plan_memory_check(split, costs)
if check.fits:
# Larger ephemeral first minimizes the running forward peak;
# ties keep chunk order, so the partition is deterministic.
order = sorted(
range(len(plans)), key=lambda i: (-costs[i].ephemeral, i)
)
return (
_SplitForwardPlan(
subforwards=tuple(plans[i] for i in order),
request_indices=tuple(tuple(chunks[i]) for i in order),
request_count=len(requests),
),
check,
)
return split, check
return None, check

def _split_rung_check(self, costs: Sequence[_SubforwardCost]) -> _MemoryCheck:
return self._memory_check_required(
sum(cost.retained for cost in costs) + max(cost.ephemeral for cost in costs)
)

@staticmethod
def _split_memory_key(plan: _SplitForwardPlan) -> bytes | None:
# Bound traversal and hashing; retain only a digest, never tensor values,
# token tables or graphs. Oversized structural identities stay unlearned.
if len(plan.subforwards) > 1024:
return None
digest, remaining, nodes = (
hashlib.sha256(),
262144 - 32 * len(plan.subforwards),
65536,
)

def feed(value: Any) -> None:
nonlocal remaining, nodes
nodes -= 1
if nodes < 0:
raise ValueError("split key node cap")
if isinstance(value, tuple):
remaining -= 2
if remaining < 0:
raise ValueError("split key byte cap")
feed(len(value))
digest.update(b"(")
for child in value:
feed(child)
digest.update(b")")
return
if value is not None and type(value) not in (str, int, bool):
raise ValueError("unsupported split key field")
if isinstance(value, str) and len(value) > 4096:
raise ValueError("split key string cap")
if isinstance(value, int) and value.bit_length() > 64:
raise ValueError("split key integer cap")
encoded = repr(value).encode()
remaining -= len(encoded) + 1
if remaining < 0:
raise ValueError("split key byte cap")
digest.update(encoded + b";")

try:
feed((plan.request_count, len(plan.subforwards)))
outer, children = digest, []
for p, indices in zip(plan.subforwards, plan.request_indices, strict=True):
digest = hashlib.sha256()
signature = p.signature
feed(
(
indices,
signature.topology,
signature.planner_coefficients,
signature.slot_group_count,
signature.request_mix,
signature.grad_enabled,
signature.grad_modes,
p.packed_tokens,
p.logical_tokens,
p.inactive_logical_tokens,
p.output_bytes,
p.output_metadata,
p.selected_max_depth,
len(p.groups),
)
)
for g in p.groups:
slot = (
None
if g.slot_ref is None
else ("checkpoint", g.slot_ref.name)
if isinstance(g.slot_ref, _LocalLoRASlotRef)
else (g.slot_ref.kind, g.slot_ref.name)
)
feed(
(
slot,
g.grad_enabled,
g.request_indices,
len(g.packed.segments),
)
)
for segment in g.packed.segments:
feed(
(
segment.sequence_indices,
segment.start,
segment.end,
segment.packed_start,
segment.group_id,
segment.parent_id,
)
)
feed(len(g.items))
for item in g.items:
feed(
(
tuple(item.input_ids.shape),
str(item.input_ids.dtype),
None
if item.labels is None
else (tuple(item.labels.shape), str(item.labels.dtype)),
item.request.top_k,
item.request.logits,
item.request.hidden_states,
)
)
children.append(digest.digest())
except ValueError:
return None
# Cost/profile updates may reorder the same children. Each digest still
# binds its original request mapping/partition; retain the observed max
# across execution orders, not an unmeasured bound for every order.
outer.update(b"".join(sorted(children)))
return outer.digest()

def _record_split_memory_floor(
self, plan: _SplitForwardPlan, baseline: int, forward_peak: int
) -> None:
# Only normal caller completion learns a floor; child retained profiles
# stay forward-only. Keep earlier child peaks despite their later resets.
key = self._split_memory_key(plan)
if key is None:
self._split_memory_floor_status = "unsupported_key"
return
if (
key not in self._split_memory_floors
and len(self._split_memory_floors) >= 1024
):
self._split_memory_floor_status = "cache_full_not_learned"
return
observed = max(
0,
max(forward_peak, int(torch.cuda.max_memory_allocated(self.device)))
- baseline,
)
self._split_memory_floors[key] = max(
self._split_memory_floors.get(key, 0), observed
)
self._split_memory_floor_status = "recorded"

def _split_plan_memory_check(
self, plan: _SplitForwardPlan, costs: Sequence[_SubforwardCost]
) -> _MemoryCheck:
required = sum(cost.retained for cost in costs) + max(
cost.ephemeral for cost in costs
)
key = self._split_memory_key(plan)
empirical = (
0
if key is None
else int(self._split_memory_floors.get(key, 0) * _MEMORY_SAFETY_FACTOR)
)
# Same existing reduction order and count; never reduce while returning
# from caller work, where a failed/empty peer may not participate.
return self._memory_check_required(max(required, empirical))

def _split_chunk_lower_cost(
self,
requests: Sequence[AnyForwardInput],
Expand Down Expand Up @@ -2649,23 +2839,6 @@ def _split_chunk_lower_cost(
logical_tokens=logical_tokens,
)
profile = self._memory_profiles.get(signature)
if profile is not None:
cap = profile.packed_tokens * _MEMORY_PROFILE_TRUST_GROWTH
if packed_tokens <= cap < unshared_packed_tokens:
# Required cost can drop when a larger layout leaves the
# profile window. Cold cost grows with packed tokens, so its
# first integer count bounds every possible post-cap layout,
# even when TP padding makes that count itself unattainable.
cold = self._subforward_cost(
packed_tokens=cap + 1,
output_bytes=output_bytes,
signature=signature,
logical_tokens=logical_tokens,
)
cost = _SubforwardCost(
required=min(cost.required, cold.required),
retained=min(cost.retained, cold.retained),
)
if (
profile is not None
and profile.retained_compute_bytes_per_token is not None
Expand Down Expand Up @@ -4839,10 +5012,10 @@ def _estimate_required_memory_bytes_from_values(
profiled_tokens = max(
packed_tokens, logical_tokens / profiled.logical_per_packed
)
if (
profiled is None
or profiled.packed_tokens * _MEMORY_PROFILE_TRUST_GROWTH < packed_tokens
):
# The trust window limits calibration growth, not the empirical floor.
# Dropping that floor beyond the window can admit a larger request that
# was refused just inside it, even below a previously observed peak.
if profiled is None:
compute = static_compute
else:
compute = max(
Expand Down
Loading
Loading