Skip to content

Report SKUAware storage reservation in planner stats - #4596

Closed
hammad45 wants to merge 1 commit into
meta-pytorch:mainfrom
hammad45:export-D116706171
Closed

Report SKUAware storage reservation in planner stats#4596
hammad45 wants to merge 1 commit into
meta-pytorch:mainfrom
hammad45:export-D116706171

Conversation

@hammad45

Copy link
Copy Markdown

Summary:
_compute_storage recognises FixedPercentageStorageReservation,
HeuristicalStorageReservation and InferenceStorageReservation, and falls
through to reserved_hbm_percent = 0.0 for anything else.
SKUAwareStorageReservation was not imported into stats.py at all, so every
successful SKUAware plan reported a reservation of zero.

Three things were wrong on each such plan:

  • Reserved Memory: HBM: 0.0 GB / 0% and Planning Memory: 100%, regardless of
    what was actually reserved.
  • The Dense Storage (per rank) and KJT Storage (per rank) blocks were
    suppressed entirely, because the same isinstance gate guards their rendering.
  • Per-rank utilization was computed against (1 - 0.0) * hbm, overstating the
    headroom the planner really had.

Observed on a dry run of an APS model: the HeuristicalStorageReservation(25%)
arm reported 23.751 GB / 25% while the SKUAware arm of the same model reported
0.0 GB / 0%.

SKUAware reserves absolute bytes anchored to a fixed home SKU, so there is no
percentage to read off it. This reports the STATIC base --
model_base_bytes when set, otherwise margin_bytes, plus
runtime_overhead_bytes -- expressed as a fraction of per-rank HBM, mirroring
the existing FixedAbsoluteStorageReservation branch. Dense and kjt are
deliberately excluded: they render as their own blocks, and including them here
would count them twice.

Differential Revision: D116706171

Summary:
`_compute_storage` recognises `FixedPercentageStorageReservation`,
`HeuristicalStorageReservation` and `InferenceStorageReservation`, and falls
through to `reserved_hbm_percent = 0.0` for anything else.
`SKUAwareStorageReservation` was not imported into `stats.py` at all, so every
successful SKUAware plan reported a reservation of zero.

Three things were wrong on each such plan:

- `Reserved Memory: HBM: 0.0 GB / 0%` and `Planning Memory: 100%`, regardless of
  what was actually reserved.
- The `Dense Storage (per rank)` and `KJT Storage (per rank)` blocks were
  suppressed entirely, because the same `isinstance` gate guards their rendering.
- Per-rank utilization was computed against `(1 - 0.0) * hbm`, overstating the
  headroom the planner really had.

Observed on a dry run of an APS model: the `HeuristicalStorageReservation(25%)`
arm reported `23.751 GB / 25%` while the SKUAware arm of the same model reported
`0.0 GB / 0%`.

SKUAware reserves absolute bytes anchored to a fixed home SKU, so there is no
percentage to read off it. This reports the STATIC base --
`model_base_bytes` when set, otherwise `margin_bytes`, plus
`runtime_overhead_bytes` -- expressed as a fraction of per-rank HBM, mirroring
the existing `FixedAbsoluteStorageReservation` branch. Dense and kjt are
deliberately excluded: they render as their own blocks, and including them here
would count them twice.

Differential Revision: D116706171
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 20, 2026
@meta-codesync

meta-codesync Bot commented Aug 20, 2026

Copy link
Copy Markdown
Contributor

@hammad45 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D116706171.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant