Merge branch 'main' into wgrad_auto_scaling

Azure · Dec 14, 2023 · 664025d · 664025d
2 parents df68631 + d562f0f
commit 664025d
Show file tree

Hide file tree

Showing 13 changed files with 640 additions and 19 deletions.
diff --git a/.gitignore b/.gitignore
@@ -138,3 +138,5 @@ dmypy.json
 
 # Cython debug symbols
 cython_debug/
+
+examples/data
diff --git a/docs/getting-started/run-msamp.md b/docs/getting-started/run-msamp.md
@@ -40,4 +40,10 @@ deepspeed cifar10_deepspeed.py --deepspeed --deepspeed_config ds_config_msamp.js
 deepspeed cifar10_deepspeed.py --deepspeed --deepspeed_config ds_config_zero_msamp.json
 ```
 
+### 4. Run cifar10 using deepspeed-ZeRO + TE with msamp enabled
+
+```bash
+deepspeed cifar10_deepspeed_te.py --deepspeed --deepspeed_config ds_config_zero_te_msamp.json
+```
+
 For more comprehensive examples, please go to [MS-AMP-Examples](https://github.com/Azure/MS-AMP-Examples).
diff --git a/docs/user-tutorial/usage.md b/docs/user-tutorial/usage.md
@@ -31,11 +31,13 @@ For enabling MS-AMP in DeepSpeed, add one line of code `from msamp import deepsp
 ```json
 "msamp": {
   "enabled": true,
-  "opt_level": "O1|O2|O3"
+  "opt_level": "O1|O2|O3",
+  "use_te": false
 }
 ```
 
 "O3" is designed for FP8 in ZeRO optimizer, so please make sure ZeRO is enabled when using "O3".
+"use_te" is designed for Transformer Engine, if you have already used Transformer Engine in your model, don't forget to set "use_te" to true.
 
 ## Usage in Megatron-DeepSpeed and Megatron-LM
Original file line number	Diff line number	Diff line change
Expand Up		@@ -138,3 +138,5 @@ dmypy.json

		# Cython debug symbols
		cython_debug/

		examples/data