Telemetry parity: adapter-contributed stats sink (extra_stats) - #4592
Open
isururanawaka wants to merge 1 commit into
Open
Telemetry parity: adapter-contributed stats sink (extra_stats)#4592isururanawaka wants to merge 1 commit into
isururanawaka wants to merge 1 commit into
Conversation
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
Contributor
|
@isururanawaka has exported this pull request. If you are a Meta employee, you can view the originating Diff in D115070614. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
D3 of the MVAI unified-planner migration: preserve MVAI's two sharding-telemetry
consumers when planning routes through
create_sharding_plan, without recompute orpolluting the neutral result contract.
MVAI's
MVAISharderStats(a torchrecStatssink) does two things when the plannerruns it: populates the in-process
model_to_sharder_infodict the variable_trainerautotuner reads, and emits the
mvai_sharding_infosignpost the MLPPS dataswarmpipeline 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 aframework Stats sink, matching the documented "stats sinks are injected as a
per-framework profile" design.
ShardingPlanAdapterBase.extra_stats()optional hook (default none), threaded ontoPlannerSessionContext.extra_statsinbuild_requestalongside the existingreport_metadata/external_trace_id/client_metadata hooks.
FbPlanReporter.build_statsappendsctx.extra_statsin both the console-only andthe manifold+scuba paths, so the framework sink runs inside the planner with the
real
StorageReservationin hand -- byte-identical, no recompute.MvaiShardingPlanAdapter.extra_stats()returnsMVAISharderStats(model_entity_id, is_invoked_by_autotuner);is_invoked_by_autotuneris a new adapter ctor param (model_entity_idalready 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