[dtensor] implement dim-0 (row) embedding sharding with MaskPartial

[wanchaol]

Stack from ghstack (oldest at bottom):

• [dtensor] add comment to clarify MaskPartial cache hit #118330
• [tp] enable rowwise embedding sharding in RowwiseParallel #118242
• -> [dtensor] implement dim-0 (row) embedding sharding with MaskPartial #118080
• [dtensor] rewrite embedding ops using op strategy #118079

This PR add support for rowwise sharded embedding by adding a
MaskPartial placement that inherits from the default partial placement,
and override the Partial constracts to construct the mask and release
the mask after the reduction

The MaskPartial placement have the potential to support other ops
sharding computation that requires a mask for semantic correctness.
currently make it live in the embedding ops but we can move it to a
common place if needed

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @fduwjj @wz337 @tianyu-l @wconstab @yf225

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/118080

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 3748220 with merge base d59c2d6 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[tianyu-l]

LGTM!
This is an elegant way of implementing row-wise embedding in DTensor. The creative use of a buffer variable in _MaskPartial slightly violates the designing principle of making Placement subclasses frozen (e.g. for caching). Nevertheless, this should be justified by the benefits it brings.

[tianyu-l]

Some remarks for what we discussed offline:

1. For _MaskPartial produced in the output of sharding propagation (either output_spec or schema_suggestions in OutputSharding), the cached self.mask_buffer.data could be filled and not released (by reductions), but still be returned with cache hit. An extremal example is that two parallel row-wise embeddings are applied on the same (replicated) input. Such cases should be rare, and if happen MaskBuffer.materialize_mask would just throw exceptions, which is OK.
2. self.mask_buffer.data is almost always not None as input to sharding propagation (because otherwise the _MaskPartial probably has been reduced to Replicate or Shard), so this is effectively forbidding cache hit when _MaskPartial is input, as is noted in a followup PR [dtensor] add comment to clarify MaskPartial cache hit #118330.

[pytorchmergebot]

@wanchaol your PR has been successfully reverted.