NVIDIA / NVIDIA/Megatron-LM

[ENHANCEMENT] Add layer name in a layer to improve code debugging

Open
#1,198 1 comment 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

**Is your feature request related to a problem? Please describe.**
I am adding new features in TranformerEngine(TE) and observe issues with model quality (gap in the loss with loss spikes).
I am debugging Megatron with TE, by storing tensor statistics in impacted layers.
But I do not have information about layer name and layer order(index) in the model topology.

**Describe the solution you'd like**
It would be great to add proper layer name with its order in the model, so that customers can use it for model debugging.

**Describe alternatives you've considered**
There are multiple frameworks which support this simple feature, e.g:
[Lingvo](https://github.com/tensorflow/lingvo/blob/c776c0f5816fc0125905b3431433c4f46d61e219/lingvo/tasks/mt/layers.py#L89) based on TF
[Praxis](https://github.com/google/praxis/blob/main/praxis/layers/transformers.py#L1741) based on JAX

**Proposed implementation**
I propose to add a layer_name filed which will be a unique name with layer hierarchy and its index/order (if there are multiple layers with the same name)

Here is an example:
'''
class TransformerBlock(MegatronModule):
"""Transformer class."""

def __init__(
self,
...
layer_name: str = "TransformerBlock",
):
# offset is implicit in TransformerLayer
self.layers = torch.nn.ModuleList(
[
build_layer(layer_spec, i + 1, f"{self.layer_name}.blocks" if self.layer_name else None)
for i, layer_spec in enumerate(self.submodules.layer_specs)
]
)
'''
**Additional context**
In our local branch, this feature is already used by multiple people.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.