Eliminate scale re-initialization #3

asfiyab-nvidia · 2023-01-05T19:07:08Z

Scale reinitialization fixed my wrapping module init in fp8_autocast and providing recipe
Update scale_factors passed to Linear, LayernormLinear, LayernormMLP layers by providing num_gemms param and providing 2 scale values per gemm.
Add configurable scale initialization for MHA and Transformer Layer

* Add TorchScript Operators * Add symbolic methods to ONNX exporter * Add tests for the ONNX export Signed-off-by: Asfiya Baig <[email protected]>

Signed-off-by: Asfiya Baig <[email protected]>

* Increase layernorm FP16 threshold * Normalize onnx file names: _ separates configs; - separates words in a single config * Add get_attn_mask_str and fix mask string * Add missing ONNX files * Moved generated ONNX files to tests/gen_onnx_models/ Signed-off-by: Asfiya Baig <[email protected]>

Signed-off-by: Asfiya Baig <[email protected]>

1. remove List import for pylint failure 2. address comments: remove state tensors from GPU 3. address comments: Update reverse_map_dtype function and add to namespace Signed-off-by: Asfiya Baig <[email protected]>

Signed-off-by: Asfiya Baig <[email protected]>

netaz · 2023-01-08T15:27:36Z

tests/test_onnx_export.py

I don't think we need this assert because it checks a TE predicate that should not affect the export process.

netaz · 2023-01-08T15:31:34Z

tests/test_onnx_export.py

You can remove num_gemms because we don't need/use it.
Please change the typing hint for scales to List.

Suggested change

def set_layer_scale(module: torch.nn.Module, scales: float, num_gemms: int=1):

def set_layer_scale(module: torch.nn.Module, scales: List[float]):

num_gemms is required for the fp8_init() call. For LayernormMLP specifically, it is set to 2. (see line below)

My bad - I missed that - ignore my comment

netaz · 2023-01-08T15:32:25Z

tests/test_onnx_export.py

Suggested change

scale_factor: list,

scale_factor: List[float],

netaz · 2023-01-08T15:32:58Z

tests/test_onnx_export.py

Suggested change

@pytest.mark.parametrize("scale_factor", [[448, 448]])

@pytest.mark.parametrize("scale_factors", [[448, 448]])

netaz · 2023-01-08T15:33:58Z

tests/test_onnx_export.py

Suggested change

@pytest.mark.parametrize("scale_factor", [[448, 448]])

@pytest.mark.parametrize("scale_factors", [[448, 448]])

Made it explicitly plural

netaz · 2023-01-08T15:37:45Z

tests/test_onnx_export.py

Suggested change

set_layer_scale(model, scale_factor, num_gemms=2)

set_layer_scale(model, scale_factors)

netaz · 2023-01-08T15:39:09Z

tests/test_onnx_export.py

Suggested change

scale_factor_qkv: list=[448, 448],

scale_factor_query: list=[112, 112],

scale_factor_kv: list=[224, 224],

scale_factor_proj: list=[448, 448]

scale_factor_qkv: List[float]=[448, 448],

scale_factor_query: List[float]=[112, 112],

scale_factor_kv: List[float]=[224, 224],

scale_factor_proj: List[float]=[448, 448]

netaz · 2023-01-08T15:41:31Z

tests/test_onnx_export.py

Suggested change

scale_factor_qkv: list,

scale_factor_query: list,

scale_factor_kv: list,

scale_factor_proj: list,

scale_factors_qkv: List[float],

scale_factors_query: List[float],

scale_factors_kv: List[float],

scale_factors_proj: List[float],

netaz · 2023-01-08T15:42:44Z

tests/test_onnx_export.py

Suggested change

scales_layernorm_mlp: list=[224, 224, 448, 448]):

scales_layernorm_mlp: List[float]=[224, 224, 448, 448]):

netaz · 2023-01-08T15:44:22Z

tests/test_onnx_export.py

Suggested change

scale_factor_qkv: list,

scale_factor_query: list,

scale_factor_kv: list,

scale_factor_proj: list,

scale_factor_layernorm_mlp: list,

scale_factor_qkv: List[float],

scale_factor_query: List[float],

scale_factor_kv: List[float],

scale_factor_proj: List[float],

scale_factor_layernorm_mlp: List[float],

Signed-off-by: Asfiya Baig <[email protected]>

1. replace variable scale_factor with scale_factors 2. Update type hints for scale_factors to be List[float] 3. Remove use of num_gemms param and add amax_history assignment Signed-off-by: Asfiya Baig <[email protected]>

asfiyab-nvidia and others added 12 commits January 4, 2023 21:32

Add ONNX export support for TE modules (#1)

f8fb17e

* Add TorchScript Operators * Add symbolic methods to ONNX exporter * Add tests for the ONNX export Signed-off-by: Asfiya Baig <[email protected]>

fixes for pylint tests

b256867

Signed-off-by: Asfiya Baig <[email protected]>

fix pylint warning in softmax.py

65231f4

Signed-off-by: Asfiya Baig <[email protected]>

move FP8 ORT lib inside tests/

2916388

Signed-off-by: Asfiya Baig <[email protected]>

enable cross attention tests

91f6c39

Signed-off-by: Asfiya Baig <[email protected]>

fix merge conflict changes

697f7e1

Signed-off-by: Asfiya Baig <[email protected]>

fix Q/DQ scale input

ebf781a

Signed-off-by: Asfiya Baig <[email protected]>

enable FP16 config when bias is disabled

b9b5477

Signed-off-by: Asfiya Baig <[email protected]>

fix pylint check errors

4812408

Signed-off-by: Asfiya Baig <[email protected]>

updates

9a7198f

1. remove List import for pylint failure 2. address comments: remove state tensors from GPU 3. address comments: Update reverse_map_dtype function and add to namespace Signed-off-by: Asfiya Baig <[email protected]>

minor fix: coding guidelines

882c462

Signed-off-by: Asfiya Baig <[email protected]>

netaz reviewed Jan 8, 2023

View reviewed changes

asfiyab-nvidia added 4 commits January 9, 2023 19:44

fix scale re-init issues

2ef1fa4

Signed-off-by: Asfiya Baig <[email protected]>

update scale arg type in definition

58920bb

Signed-off-by: Asfiya Baig <[email protected]>

add scales to mha and transformer layer submodules

486db6b

Signed-off-by: Asfiya Baig <[email protected]>

address comments

e83246e

1. replace variable scale_factor with scale_factors 2. Update type hints for scale_factors to be List[float] 3. Remove use of num_gemms param and add amax_history assignment Signed-off-by: Asfiya Baig <[email protected]>

asfiyab-nvidia force-pushed the dev-fix-scale-init branch from 325fe62 to e83246e Compare January 9, 2023 20:11

asfiyab-nvidia force-pushed the dev-onnx-export-support branch from 119a0ec to ab4410f Compare January 17, 2023 20:35

	def set_layer_scale(module: torch.nn.Module, scales: float, num_gemms: int=1):
	def set_layer_scale(module: torch.nn.Module, scales: List[float]):

	@pytest.mark.parametrize("scale_factor", [[448, 448]])
	@pytest.mark.parametrize("scale_factors", [[448, 448]])

	set_layer_scale(model, scale_factor, num_gemms=2)
	set_layer_scale(model, scale_factors)

-    scale_factor_qkv: list=[448, 448],
-    scale_factor_query: list=[112, 112],
-    scale_factor_kv: list=[224, 224],
-    scale_factor_proj: list=[448, 448]
+    scale_factor_qkv: List[float]=[448, 448],
+    scale_factor_query: List[float]=[112, 112],
+    scale_factor_kv: List[float]=[224, 224],
+    scale_factor_proj: List[float]=[448, 448]

-    scale_factor_qkv: list,
-    scale_factor_query: list,
-    scale_factor_kv: list,
-    scale_factor_proj: list,
+    scale_factors_qkv: List[float],
+    scale_factors_query: List[float],
+    scale_factors_kv: List[float],
+    scale_factors_proj: List[float],

	scales_layernorm_mlp: list=[224, 224, 448, 448]):
	scales_layernorm_mlp: List[float]=[224, 224, 448, 448]):

-    scale_factor_qkv: list,
-    scale_factor_query: list,
-    scale_factor_kv: list,
-    scale_factor_proj: list,
-    scale_factor_layernorm_mlp: list,
+    scale_factor_qkv: List[float],
+    scale_factor_query: List[float],
+    scale_factor_kv: List[float],
+    scale_factor_proj: List[float],
+    scale_factor_layernorm_mlp: List[float],

Eliminate scale re-initialization #3

Are you sure you want to change the base?

Eliminate scale re-initialization #3

Uh oh!

Conversation

asfiyab-nvidia commented Jan 5, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants