[ENHANCEMENT] Add layer name in a layer to improve code debugging
- 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
Assessment
This issue has not been assessed yet.