Skip to content

Telemetry parity: adapter-contributed stats sink (extra_stats) - #4592

Open
isururanawaka wants to merge 1 commit into
meta-pytorch:mainfrom
isururanawaka:export-D115070614
Open

Telemetry parity: adapter-contributed stats sink (extra_stats)#4592
isururanawaka wants to merge 1 commit into
meta-pytorch:mainfrom
isururanawaka:export-D115070614

Conversation

@isururanawaka

Copy link
Copy Markdown
Contributor

Summary:
D3 of the MVAI unified-planner migration: preserve MVAI's two sharding-telemetry
consumers when planning routes through create_sharding_plan, without recompute or
polluting the neutral result contract.

MVAI's MVAISharderStats (a torchrec Stats sink) does two things when the planner
runs it: populates the in-process model_to_sharder_info dict the variable_trainer
autotuner reads, and emits the mvai_sharding_info signpost the MLPPS dataswarm
pipeline consumes. On the unified path the reporter builds its own sink list, so
MVAISharderStats never ran and both consumers would break.

Clean fix (chosen over recomputing per-rank storage_hbm from the result -- which
would need dense/kjt exposed on ShardingPlanResult): let the adapter contribute a
framework Stats sink, matching the documented "stats sinks are injected as a
per-framework profile" design.

  • ShardingPlanAdapterBase.extra_stats() optional hook (default none), threaded onto
    PlannerSessionContext.extra_stats in build_request alongside the existing
    report_metadata/external_trace_id/client_metadata hooks.
  • FbPlanReporter.build_stats appends ctx.extra_stats in both the console-only and
    the manifold+scuba paths, so the framework sink runs inside the planner with the
    real StorageReservation in hand -- byte-identical, no recompute.
  • MvaiShardingPlanAdapter.extra_stats() returns MVAISharderStats(model_entity_id, is_invoked_by_autotuner); is_invoked_by_autotuner is a new adapter ctor param (model_entity_id already existed on the adapter).

Result: the autotuner handoff and the MLPPS signpost keep working unchanged; the
neutral result/config contracts stay generic.

Differential Revision: D115070614

Summary:
D3 of the MVAI unified-planner migration: preserve MVAI's two sharding-telemetry
consumers when planning routes through `create_sharding_plan`, without recompute or
polluting the neutral result contract.

MVAI's `MVAISharderStats` (a torchrec `Stats` sink) does two things when the planner
runs it: populates the in-process `model_to_sharder_info` dict the variable_trainer
autotuner reads, and emits the `mvai_sharding_info` signpost the MLPPS dataswarm
pipeline consumes. On the unified path the reporter builds its own sink list, so
MVAISharderStats never ran and both consumers would break.

Clean fix (chosen over recomputing per-rank storage_hbm from the result -- which
would need dense/kjt exposed on `ShardingPlanResult`): let the adapter contribute a
framework Stats sink, matching the documented "stats sinks are injected as a
per-framework profile" design.
- `ShardingPlanAdapterBase.extra_stats()` optional hook (default none), threaded onto
  `PlannerSessionContext.extra_stats` in `build_request` alongside the existing
  report_metadata/external_trace_id/client_metadata hooks.
- `FbPlanReporter.build_stats` appends `ctx.extra_stats` in both the console-only and
  the manifold+scuba paths, so the framework sink runs inside the planner with the
  real `StorageReservation` in hand -- byte-identical, no recompute.
- `MvaiShardingPlanAdapter.extra_stats()` returns `MVAISharderStats(model_entity_id,
  is_invoked_by_autotuner)`; `is_invoked_by_autotuner` is a new adapter ctor param (`model_entity_id` already existed on the adapter).

Result: the autotuner handoff and the MLPPS signpost keep working unchanged; the
neutral result/config contracts stay generic.

Differential Revision: D115070614
@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

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

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