Skip to content

Replace per-rank table_names dump with a fingerprint - #4540

Open
kaanbaloglu wants to merge 1 commit into
meta-pytorch:mainfrom
kaanbaloglu:export-D114819716
Open

Replace per-rank table_names dump with a fingerprint#4540
kaanbaloglu wants to merge 1 commit into
meta-pytorch:mainfrom
kaanbaloglu:export-D114819716

Conversation

@kaanbaloglu

Copy link
Copy Markdown
Contributor

Summary:
This diff replaces the per-rank sharding-plan dump in the two sharder events with a small fingerprint, and keeps the full map on one rank per sharding group.

ShardedEmbeddingBagCollection.__init__ logs the full table-name to sharding-type map on every rank. Every rank builds an identical copy, so a large job repeats the same multi-kilobyte payload once per rank. That is a tiny share of the event stream's rows and most of its bytes. The planner already persists the same plan durably elsewhere. ShardedEmbeddingCollection does the same thing.

Each sharder now emits two events. Every rank emits <Module>.sharding_plan_fingerprint with num_tables, per-sharding-type counts, and a fingerprint. One rank per sharding group also emits the existing <Module>.table_names, unchanged. The fingerprint comes from sharding_plan_fingerprint in torchrec/distributed/utils.py. It uses blake2b, not hash(), which is salted per process and not comparable across ranks. Every rank keeps a fingerprint because some jobs shard heterogeneously, so a payload can sit on a rank band that excludes rank 0. The gate is is_plan_leader. It keys on rank 0 of the env's own process group. The plan is broadcast over that group, so it is identical across it and can differ across groups. Under 2D sharding that group is global_pg. Tower and pipeline-stage sharding build an env per subgroup, and most of those subgroups do not contain global rank 0. A job-wide gate would drop their plans. One emission per process group removes every duplicate copy and keeps every distinct plan.

Differential Revision: D114819716

Summary:
This diff replaces the per-rank sharding-plan dump in the two sharder events with a small fingerprint, and keeps the full map on one rank per sharding group.

`ShardedEmbeddingBagCollection.__init__` logs the full table-name to sharding-type map on every rank. Every rank builds an identical copy, so a large job repeats the same multi-kilobyte payload once per rank. That is a tiny share of the event stream's rows and most of its bytes. The planner already persists the same plan durably elsewhere. `ShardedEmbeddingCollection` does the same thing.

Each sharder now emits two events. Every rank emits `<Module>.sharding_plan_fingerprint` with `num_tables`, per-sharding-type counts, and a fingerprint. One rank per sharding group also emits the existing `<Module>.table_names`, unchanged. The fingerprint comes from `sharding_plan_fingerprint` in `torchrec/distributed/utils.py`. It uses `blake2b`, not `hash()`, which is salted per process and not comparable across ranks. Every rank keeps a fingerprint because some jobs shard heterogeneously, so a payload can sit on a rank band that excludes rank 0. The gate is `is_plan_leader`. It keys on rank 0 of the env's own process group. The plan is broadcast over that group, so it is identical across it and can differ across groups. Under 2D sharding that group is `global_pg`. Tower and pipeline-stage sharding build an env per subgroup, and most of those subgroups do not contain global rank 0. A job-wide gate would drop their plans. One emission per process group removes every duplicate copy and keeps every distinct plan.

Differential Revision: D114819716
@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 13, 2026
@meta-codesync

meta-codesync Bot commented Aug 13, 2026

Copy link
Copy Markdown
Contributor

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

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