Lightning-AI / Lightning-AI/pytorch-lightning
Distributed group information for MOE layer
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
📚 Documentation
Thank you for maintaining this amazing repository.
I am integrating MOE layers into my model architecture, which I am training using lightning.
I am using megablocks implementation due to its wider adoption. One of the variables required to enable moe_expert_model_parallelism is distributed group information(https://github.com/databricks/megablocks/blob/main/megablocks/layers/memory_test.py#L97C5-L97C10). I am wondering if there is a way to access this information in LightningModule before model initialization.
I would appreciate any guidance you can provide on how to access the group variable, even if it is not straightforward with the current lightning API. Thank you very much for your time and help!
Regards,
Pradyumna.
cc @lantiga @borda
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the LightningModule initialization lifecycle and the distributed setup APIs, then compare them with the distributed group usage shown in megablocks/layers/memory_test.py. Determine whether the group information is available before model initialization and identify the smallest documented or API-level change needed to make that access possible.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100