TransformerEngine Integration #1282

aurelion-source · 2024-09-16T00:28:04Z

Adds:

TELinear
TELayerNormMLP
TEColumnParallelLinear
TERowParallelLinear
TEMultiheadAttention
TEDelayedScaling (in progress)

CLAassistant · 2024-09-16T00:28:10Z

All committers have signed the CLA.

- Implemented delayed scaling

megatron/model/utils.py

- Implemented delayed scaling

…forward

megatron/model/transformer_engine.py

Quentin-Anthony · 2024-10-08T20:32:42Z

@aurelion-source

I tried to run this with an NGC container (specifically nvcr.io/nvidia/pytorch:23.10-py3)

Traceback (most recent call last):                                                                                                                                                                                                     
  File "/workspace/gpt-neox-nawras/train.py", line 35, in <module>                                                                                                                                                                     
    main()                                                                                                                                                                                                                             
  File "/workspace/gpt-neox-nawras/train.py", line 31, in main                                                                                                                                                                         
    pretrain(neox_args=neox_args)                                                                                                                                                                                                      
  File "/workspace/gpt-neox-nawras/megatron/training.py", line 251, in pretrain                                                                                                                                                        
    model, optimizer, lr_scheduler, reference_model = setup_model_and_optimizer(                                                                                                                                                       
  File "/workspace/gpt-neox-nawras/megatron/training.py", line 1153, in setup_model_and_optimizer                                                                                                                                      
    model = get_model(neox_args=neox_args, use_cache=use_cache)                                                                                                                                                                        
  File "/workspace/gpt-neox-nawras/megatron/training.py", line 882, in get_model                                                                                                                                                       
    model = GPT2ModelPipe(                                                                                                                                                                                                             
  File "/workspace/gpt-neox-nawras/megatron/model/gpt2_model.py", line 131, in __init__
    super().__init__(                                                                                                                                                                                                                  
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 212, in __init__
    self._build()
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 268, in _build
    module = layer.build()                                                                                         
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 74, in build                                                                                                                                   
    return self.typename(*self.module_args, **self.module_kwargs)
  File "/workspace/gpt-neox-nawras/megatron/model/transformer.py", line 1030, in __init__
    from megatron.model.transformer_engine import TEMultiheadAttention                
  File "/workspace/gpt-neox-nawras/megatron/model/transformer_engine.py", line 97, in <module>                                                                                                                                         
    class TELinear(te.pytorch.Linear):                                                                             
AttributeError: module 'transformer_engine' has no attribute 'pytorch'

Updating the import to import transformer_engine.pytorch as te updates my error to:

Traceback (most recent call last):                                                                                                                                                                                                       File "/workspace/gpt-neox-nawras/train.py", line 35, in <module>                                                                                                                                                                         main()                                                                                                                                                                                                                               File "/workspace/gpt-neox-nawras/train.py", line 31, in main                                                                                                                                                                             pretrain(neox_args=neox_args)                                                                                                                                                                                                        File "/workspace/gpt-neox-nawras/megatron/training.py", line 251, in pretrain                                                                                                                                                            model, optimizer, lr_scheduler, reference_model = setup_model_and_optimizer(                                                                                                                                                         File "/workspace/gpt-neox-nawras/megatron/training.py", line 1153, in setup_model_and_optimizer                                                                                                                                          model = get_model(neox_args=neox_args, use_cache=use_cache)                                                                                                                                                                          File "/workspace/gpt-neox-nawras/megatron/training.py", line 882, in get_model                                                                                                                                                       
    model = GPT2ModelPipe(                          
  File "/workspace/gpt-neox-nawras/megatron/model/gpt2_model.py", line 131, in __init__
    super().__init__(                           
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 212, in __init__
    self._build()                               
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 268, in _build
    module = layer.build()                      
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 74, in build
    return self.typename(*self.module_args, **self.module_kwargs)
  File "/workspace/gpt-neox-nawras/megatron/model/transformer.py", line 1030, in __init__
    from megatron.model.transformer_engine import TEMultiheadAttention
  File "/workspace/gpt-neox-nawras/megatron/model/transformer_engine.py", line 45, in <module>
    import transformer_engine.pytorch as te
  File "/usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/__init__.py", line 63, in <module>
    _load_library()
  File "/usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/__init__.py", line 55, in _load_library
    so_path = next(so_dir.glob(f"{module_name}.*.{extension}"))
StopIteration

Basically the same as NVIDIA/TransformerEngine#1014. It's a version issue with TE 1.9.0 and 1.10.0. It would be preferable to allow these versions to function so that NGC containers and pip installs can be used.

Even using pip install -r requirements/requirements-transformerengine.txt for latest stable source-build fails with:

Traceback (most recent call last):                                                                                                                                                                                                       File "/workspace/gpt-neox-nawras/train.py", line 35, in <module>                                                                                                                                                                         main()                                                                                                                                                                                                                               File "/workspace/gpt-neox-nawras/train.py", line 31, in main                                                                                                                                                                             pretrain(neox_args=neox_args)                                                                                                                                                                                                        File "/workspace/gpt-neox-nawras/megatron/training.py", line 251, in pretrain                                                                                                                                                            model, optimizer, lr_scheduler, reference_model = setup_model_and_optimizer(                                                                                                                                                         File "/workspace/gpt-neox-nawras/megatron/training.py", line 1153, in setup_model_and_optimizer                                                                                                                                          model = get_model(neox_args=neox_args, use_cache=use_cache)                                                                                                                                                                        
  File "/workspace/gpt-neox-nawras/megatron/training.py", line 882, in get_model
    model = GPT2ModelPipe(                      
  File "/workspace/gpt-neox-nawras/megatron/model/gpt2_model.py", line 131, in __init__
    super().__init__(                                                                                              
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 212, in __init__
    self._build()
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 268, in _build
    module = layer.build()
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 74, in build
    return self.typename(*self.module_args, **self.module_kwargs)
  File "/workspace/gpt-neox-nawras/megatron/model/transformer.py", line 1030, in __init__
    from megatron.model.transformer_engine import TEMultiheadAttention
  File "/workspace/gpt-neox-nawras/megatron/model/transformer_engine.py", line 97, in <module>
    class TELinear(te.pytorch.Linear):
AttributeError: module 'transformer_engine.pytorch' has no attribute 'pytorch'. Did you mean: 'torch'?

aurelion-source · 2024-10-16T21:31:33Z

@aurelion-source

I tried to run this with an NGC container (specifically nvcr.io/nvidia/pytorch:23.10-py3)

Traceback (most recent call last):                                                                                                                                                                                                     
  File "/workspace/gpt-neox-nawras/train.py", line 35, in <module>                                                                                                                                                                     
    main()                                                                                                                                                                                                                             
  File "/workspace/gpt-neox-nawras/train.py", line 31, in main                                                                                                                                                                         
    pretrain(neox_args=neox_args)                                                                                                                                                                                                      
  File "/workspace/gpt-neox-nawras/megatron/training.py", line 251, in pretrain                                                                                                                                                        
    model, optimizer, lr_scheduler, reference_model = setup_model_and_optimizer(                                                                                                                                                       
  File "/workspace/gpt-neox-nawras/megatron/training.py", line 1153, in setup_model_and_optimizer                                                                                                                                      
    model = get_model(neox_args=neox_args, use_cache=use_cache)                                                                                                                                                                        
  File "/workspace/gpt-neox-nawras/megatron/training.py", line 882, in get_model                                                                                                                                                       
    model = GPT2ModelPipe(                                                                                                                                                                                                             
  File "/workspace/gpt-neox-nawras/megatron/model/gpt2_model.py", line 131, in __init__
    super().__init__(                                                                                                                                                                                                                  
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 212, in __init__
    self._build()
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 268, in _build
    module = layer.build()                                                                                         
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 74, in build                                                                                                                                   
    return self.typename(*self.module_args, **self.module_kwargs)
  File "/workspace/gpt-neox-nawras/megatron/model/transformer.py", line 1030, in __init__
    from megatron.model.transformer_engine import TEMultiheadAttention                
  File "/workspace/gpt-neox-nawras/megatron/model/transformer_engine.py", line 97, in <module>                                                                                                                                         
    class TELinear(te.pytorch.Linear):                                                                             
AttributeError: module 'transformer_engine' has no attribute 'pytorch'

Updating the import to import transformer_engine.pytorch as te updates my error to:

Traceback (most recent call last):                                                                                                                                                                                                       File "/workspace/gpt-neox-nawras/train.py", line 35, in <module>                                                                                                                                                                         main()                                                                                                                                                                                                                               File "/workspace/gpt-neox-nawras/train.py", line 31, in main                                                                                                                                                                             pretrain(neox_args=neox_args)                                                                                                                                                                                                        File "/workspace/gpt-neox-nawras/megatron/training.py", line 251, in pretrain                                                                                                                                                            model, optimizer, lr_scheduler, reference_model = setup_model_and_optimizer(                                                                                                                                                         File "/workspace/gpt-neox-nawras/megatron/training.py", line 1153, in setup_model_and_optimizer                                                                                                                                          model = get_model(neox_args=neox_args, use_cache=use_cache)                                                                                                                                                                          File "/workspace/gpt-neox-nawras/megatron/training.py", line 882, in get_model                                                                                                                                                       
    model = GPT2ModelPipe(                          
  File "/workspace/gpt-neox-nawras/megatron/model/gpt2_model.py", line 131, in __init__
    super().__init__(                           
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 212, in __init__
    self._build()                               
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 268, in _build
    module = layer.build()                      
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 74, in build
    return self.typename(*self.module_args, **self.module_kwargs)
  File "/workspace/gpt-neox-nawras/megatron/model/transformer.py", line 1030, in __init__
    from megatron.model.transformer_engine import TEMultiheadAttention
  File "/workspace/gpt-neox-nawras/megatron/model/transformer_engine.py", line 45, in <module>
    import transformer_engine.pytorch as te
  File "/usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/__init__.py", line 63, in <module>
    _load_library()
  File "/usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/__init__.py", line 55, in _load_library
    so_path = next(so_dir.glob(f"{module_name}.*.{extension}"))
StopIteration

Basically the same as NVIDIA/TransformerEngine#1014. It's a version issue with TE 1.9.0 and 1.10.0. It would be preferable to allow these versions to function so that NGC containers and pip installs can be used.

Even using pip install -r requirements/requirements-transformerengine.txt for latest stable source-build fails with:

Traceback (most recent call last):                                                                                                                                                                                                       File "/workspace/gpt-neox-nawras/train.py", line 35, in <module>                                                                                                                                                                         main()                                                                                                                                                                                                                               File "/workspace/gpt-neox-nawras/train.py", line 31, in main                                                                                                                                                                             pretrain(neox_args=neox_args)                                                                                                                                                                                                        File "/workspace/gpt-neox-nawras/megatron/training.py", line 251, in pretrain                                                                                                                                                            model, optimizer, lr_scheduler, reference_model = setup_model_and_optimizer(                                                                                                                                                         File "/workspace/gpt-neox-nawras/megatron/training.py", line 1153, in setup_model_and_optimizer                                                                                                                                          model = get_model(neox_args=neox_args, use_cache=use_cache)                                                                                                                                                                        
  File "/workspace/gpt-neox-nawras/megatron/training.py", line 882, in get_model
    model = GPT2ModelPipe(                      
  File "/workspace/gpt-neox-nawras/megatron/model/gpt2_model.py", line 131, in __init__
    super().__init__(                                                                                              
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 212, in __init__
    self._build()
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 268, in _build
    module = layer.build()
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/pipe/module.py", line 74, in build
    return self.typename(*self.module_args, **self.module_kwargs)
  File "/workspace/gpt-neox-nawras/megatron/model/transformer.py", line 1030, in __init__
    from megatron.model.transformer_engine import TEMultiheadAttention
  File "/workspace/gpt-neox-nawras/megatron/model/transformer_engine.py", line 97, in <module>
    class TELinear(te.pytorch.Linear):
AttributeError: module 'transformer_engine.pytorch' has no attribute 'pytorch'. Did you mean: 'torch'?

@Quentin-Anthony
Based on NVIDIA/TransformerEngine#1014 (comment), the issue is due to TE not detecting PyTorch during the build process. It's not a version issue.

The suggest solution NVTE_FRAMEWORK=pytorch pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable works fine. However, this requires setting an environment variable before the pip command.

Alternatively, pip install transformer-engine[pytorch] forces TE to build with PyTorch support as well.
I've updated and tested ./requirements/requirements-transformer-engine.txt accordingly.

aurelion-source added 7 commits September 14, 2024 23:41

Implemented ColumnParallelLinear with Transformer-Engine

5382c8a

Implemented RowParallelLinear with Transformer-Engine

fa887b7

Implemented LayerNormMLP with Transformer-Engine

0a6f140

Implemented MultiheadAttention with Transformer-Engine

5cba717

Cleaned up transformer.py

94e552c

Cleaned up neox_args

40e1019

Cleaned up neox_args

885e72c

aurelion-source self-assigned this Sep 16, 2024

aurelion-source requested a review from Quentin-Anthony as a code owner September 16, 2024 00:28

aurelion-source added 2 commits September 26, 2024 08:18

- Fixed TE_MHA and added rope support

fe8f22a

- Implemented delayed scaling

Fixed mixed files.

ee42a31

Quentin-Anthony reviewed Sep 26, 2024

View reviewed changes

megatron/model/utils.py Outdated Show resolved Hide resolved

aurelion-source added 10 commits September 27, 2024 15:12

Implemented ColumnParallelLinear with Transformer-Engine

8961dd7

Implemented RowParallelLinear with Transformer-Engine

5162d54

Implemented LayerNormMLP with Transformer-Engine

3cad89c

Implemented MultiheadAttention with Transformer-Engine

eedb6c2

Cleaned up transformer.py

36ad680

Cleaned up neox_args

6963103

Cleaned up neox_args

afc9c92

- Fixed TE_MHA and added rope support

a0e7acd

- Implemented delayed scaling

Fixed mixed files.

0b4bdc5

Merge branch 'te' of github.com:aurelion-source/gpt-neox into te

d559be9

aurelion-source closed this Sep 27, 2024

aurelion-source reopened this Sep 27, 2024

Changed get_linear name

bb76510

tf-nv mentioned this pull request Sep 30, 2024

TE integration via full TransformerLayer #1297

Open

Added rng tracker to lnmlp and placed rope in te_mha init instead of …

43cf4ee

…forward

Quentin-Anthony reviewed Oct 1, 2024

View reviewed changes

megatron/model/transformer_engine.py Show resolved Hide resolved

Updated fp8 arguments to te_fp8

42716f2

aurelion-source and others added 4 commits October 1, 2024 05:28

Added EAI copyright

b3255e6

Merge branch 'main' into te

afeff03

precommit

98f0388

add sample TE config

7e7dbfb

Quentin-Anthony and others added 3 commits October 8, 2024 13:32

add te to readme

5757be6

remove pip install prefix from reqs file

9ea3dcf

Force TE pytorch in requirements file

f3e40e9

Merge remote-tracking branch 'upstream/main' into te

7bf9c44

Quentin-Anthony approved these changes Dec 19, 2024

View reviewed changes

Quentin-Anthony merged commit 8900d05 into EleutherAI:main Dec 19, 2024
1 of 2 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

TransformerEngine Integration #1282

TransformerEngine Integration #1282

aurelion-source commented Sep 16, 2024

CLAassistant commented Sep 16, 2024 •

edited

Loading

Quentin-Anthony commented Oct 8, 2024

aurelion-source commented Oct 16, 2024

TransformerEngine Integration #1282

TransformerEngine Integration #1282

Conversation

aurelion-source commented Sep 16, 2024

CLAassistant commented Sep 16, 2024 • edited Loading

Quentin-Anthony commented Oct 8, 2024

aurelion-source commented Oct 16, 2024

CLAassistant commented Sep 16, 2024 •

edited

Loading