feat: Reduce reference model memory with with parallel logprob computation#608

Merged

felipemello1 merged 20 commits intometa-pytorch:mainfrom

gitlost-murali:optimize-ref-model-usage

Dec 9, 2025

Commits on Dec 8, 2025

feat: Add parallel log-prob computation for vocab-sharded tensors in reference model
gitlost-murali
committed
refactor: use util function
gitlost-murali
committed
refactor: remove redundant code and reuse compute_logprobs function
gitlost-murali
committed
test: add test to verify correctness with compute_logprobs
gitlost-murali
committed
chore: make routing explicit
gitlost-murali
committed
style: clean up inline comments in parallel_logprobs
gitlost-murali
committed
refactor: move compute_logprobs_parallel into ops alongside compute_logprobs
gitlost-murali
committed
refactor: simplify and refactor logprobs parallel method
gitlost-murali
committed
test: move parallel logprobs test to test_ops.py reflecting folder structure
gitlost-murali
committed
chore: merge single use declarations
gitlost-murali
committed
fix: safely handle edgecase of uneven vocab shards in tensor-parallel logprobs
gitlost-murali
committed
fix: fix compute_logprobs_parallel import
gitlost-murali
committed
feat: add torch compile to further reduce the reference model GPU usage in non-sharded and sharded computation
gitlost-murali
committed
refactor: Simplify sharded logprobs computation using loss_parallel context
gitlost-murali
committed
fix: convert DTensor output to regular tensor after loss_parallel
gitlost-murali
committed
refactor: make compile configurable for logprobs computation
gitlost-murali
committed
refactor: remove redundant logging steps as on main
gitlost-murali
committed

Commits on Dec 9, 2025