[Determinism] determinism code clean up
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
There is some remaining megatron code that diverges between deterministic and non-deterministic. I think we should try to clean it up. Notes follow.
- [layers.py](https://github.com/NVIDIA/Megatron-LM/blob/d347eef80/megatron/core/tensor_parallel/layers.py#L351-L354) – Claims that `F.embedding` has a "non-deterministic backward function". We should try to remove this case and determine one way or the other.
- [determinism.py](https://github.com/NVIDIA/Megatron-LM/blob/d347eef80/megatron/core/ssm/ops/common/determinism.py#L106-L124) – `alloc_tile_workspace()` and `finalize_tile_workspace()` are zeroing out the workspace unnecessarily. Similar comments apply to the same code in the mamba repository.
- [gdn.py](https://github.com/NVIDIA/Megatron-LM/blob/d347eef80/megatron/core/ssm/gated_delta_net/gdn.py#L58-L61) – This is choosing between FLA and torch native version of the `gated_delta_rule` because the FLA code is non-deterministic. If the [FLA code](https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/gated_delta_rule/chunk.py) is non-deterministic, I think we could fix it.
- [moe_utils.py](https://github.com/NVIDIA/Megatron-LM/blob/d347eef80/megatron/core/transformer/moe/moe_utils.py#L600-L617) – Is the non-deterministic code path any faster? The deterministic path also contains a redundant call to `torch.zeros()` which should be removed.
- [coordinator.py](https://github.com/NVIDIA/Megatron-LM/blob/d347eef80/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py#L182-L185) – If reading correctly, we give up deterministic inference to avoid a sort at initialization time? Is this really worth it?
- [inference permute kernels](https://github.com/NVIDIA/Megatron-LM/blob/d347eef80/megatron/core/inference/moe/permute.py#L493-L501) – there has been a massive change along mainline due to this [PR](https://github.com/NVIDIA/Megatron-LM/pull/4871) by Peter Dykas (make megatron batch invariant).
Contributor guide
Research direction
Start with the referenced sections in layers.py, ssm/ops/common/determinism.py, gdn.py, moe_utils.py, coordinator.py, and inference/moe/permute.py, then compare the deterministic and non-deterministic paths and the linked batch-invariance change. Done means each listed discrepancy has an evidence-based resolution, with unnecessary workspace or tensor initialization removed only where safe and determinism behavior preserved.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, performance
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100