NVIDIA / NVIDIA/Megatron-LM

[Determinism] determinism code clean up

Open
#6,547 3 comments 0 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.