Port output distribution classes to DynamicEmb

[z52527]

Description

#296
Port TorchRec's output distribution classes to DynamicEmb library, enabling future performance optimizations.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[z52527]

Track CI here.

[greptile-apps]

Greptile Overview

Greptile Summary

• Adds DynamicEmb equivalents of TorchRec RW output distribution modules for sequence (AlltoAll) and pooled (reduce-scatter) embeddings in output_dist.py.
• Hooks these output distribution modules into DynamicEmb’s RW sharding implementations in planner/rw_sharding.py via new create_output_dist() overrides.
• Enables future performance work by keeping distribution logic local to DynamicEmb while still leveraging TorchRec comm primitives and qcomm codec registry integration.

Confidence Score: 3/5

• This PR is mergeable after fixing two concrete runtime-crash paths in the new output distribution wiring.
• The changes are localized and follow TorchRec patterns, but (1) the sharding code can pass a None process group into dist constructors (leading to immediate crashes), and (2) pooled output dist can be initialized in variable-batch mode and later invoked without required ctx args, also crashing.
• corelib/dynamicemb/dynamicemb/output_dist.py, corelib/dynamicemb/dynamicemb/planner/rw_sharding.py

Important Files Changed

Filename	Overview
corelib/dynamicemb/dynamicemb/output_dist.py	Adds RW output distribution modules for sequence (AlltoAll) and pooled (reduce-scatter) embeddings; there is a runtime hazard if the first call initializes variable-batch reduce-scatter and later calls omit sharding_ctx.
corelib/dynamicemb/dynamicemb/planner/rw_sharding.py	Wires new output distribution classes into RW dynamic embedding shardings; currently passes Optional process group into dist constructors (can crash when pg is None) and leaves TODO/commented uncertainty.

Sequence Diagram

sequenceDiagram
  participant Planner as rw_sharding.py
  participant OutDist as output_dist.py
  participant A2A as SequenceEmbeddingsAllToAll
  participant RS as PooledEmbeddingsReduceScatter
  participant VRS as VariableBatchPooledEmbeddingsReduceScatter

  Planner->>OutDist: create_output_dist(pg, num_features/embedding_dims, codecs)
  Note over Planner,OutDist: pg is currently Optional in sharding classes

  alt Sequence embeddings
    Planner->>OutDist: RwSequenceEmbeddingDist.forward(local_embs, SequenceShardingContext)
    OutDist->>A2A: alltoall(lengths/input_splits/output_splits/...)
    A2A-->>OutDist: Awaitable[Tensor]
  else Pooled embeddings
    Planner->>OutDist: RwPooledEmbeddingDist.forward(local_embs, EmbeddingShardingContext?)
    alt first call initializes dist module
      OutDist->>RS: init reduce-scatter
      OutDist->>VRS: init variable-batch reduce-scatter (if ctx.variable_batch_per_feature)
    end
    alt ctx is None
      OutDist->>RS: reduce_scatter(local_embs)
    else ctx.variable_batch_per_feature
      OutDist->>VRS: reduce_scatter(local_embs, batch_size_per_rank_per_feature, embedding_dims)
    else ctx provided
      OutDist->>RS: reduce_scatter(local_embs, input_splits=batch_size_per_rank)
    end
  end

Loading

[greptile-apps]

_{2 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Passes Optional ProcessGroup

create_output_dist() passes self._pg through to RwSequenceEmbeddingDist, but self._pg is typed as Optional[ProcessGroup] (and you’ve added a pyre-fixme to silence it). If self._pg is actually None at runtime (e.g., in non-distributed / single-rank setups), RwSequenceEmbeddingDist.__init__ will call pg.size() and crash. This needs a real guard or to ensure _pg is always non-None before constructing the dist module (same pattern in pooled sharding too).

Also appears in: corelib/dynamicemb/dynamicemb/planner/rw_sharding.py:264-274.

[greptile-apps]

Missing context for init

When sharding_ctx is None, forward() returns PooledEmbeddingsReduceScatter(local_embs) (line 134-135) but _dist may have been created with a variable-batch module based on the first call’s sharding_ctx (line 131-133 / 151-155). If the first invocation had variable_batch_per_feature=True and a later call passes sharding_ctx=None, this will call a VariableBatchPooledEmbeddingsReduceScatter without its required args, causing a runtime error. Either require sharding_ctx always be provided, or make _dist selection independent of the first call.