feat: support group_norm, batch_norm, and layer_norm #2330

zewenli98 · 2023-09-20T23:21:24Z

Description

Update batch_norm and layer_norm

Fixes #2225

Type of change

Bug fix (non-breaking change which fixes an issue)

Checklist:

My code follows the style guidelines of this project (You can use the linters)
I have performed a self-review of my own code
I have commented my code, particularly in hard-to-understand areas and hacks
I have made corresponding changes to the documentation
I have added tests to verify my fix or my feature
New and existing unit tests pass locally with my changes
I have added the relevant labels to my PR in so that relevant reviewers are notified

gs-olive

Updates look great - added some suggestions to better follow the Torch schemas for these functions

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py

gs-olive · 2023-09-23T01:06:10Z

py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py

+    if weight is None:
+        weight = np.array(1.0)
+
+    if bias is None:
+        bias = np.array(0.0)
+
+    if running_mean is None:
+        running_mean = np.array(0.0)
+
+    if running_var is None:
+        running_var = np.array(1.0)


For these, it should be okay to not cast to np.array in the converter (instead leave them as ints or floats), since to_numpy should dictate this casting behavior for ints and floats. Specifically, one small difference is that I think np.array(1.0) has shape () (0D), but to_numpy generally adds a dimension, to make it 1D.

gs-olive · 2023-09-23T01:06:40Z

py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py

+    if weight is None:
+        weight = np.array(1.0)
+
+    if bias is None:
+        bias = np.array(0.0)


See above comment

gs-olive · 2023-09-23T01:06:53Z

py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py

+    if weight is None:
+        weight = np.array(1.0)
+
+    if bias is None:
+        bias = np.array(0.0)


See above comment

Since line 189 is shape = weight.shape and lines 191 and 192 call weight.reshape and bias.reshape, I think weight and bias shouldn't be scalars.

I see - in that case, it might be preferable to use to_numpy(0.0), for instance, to get back a default-formatted numpy array for the float default. Additionally, I noticed the code below has some issues:

gamma = to_numpy(weight.reshape(*shape)) ##### Above is invalid, since the reshape should apply to the numpy output. It should instead be: gamma = to_numpy(weight).reshape(shape)

The same as the above applies for beta.

Additionally, lines 194 - 196 should be using get_axes_for_reduce_op, as here:

TensorRT/py/torch_tensorrt/dynamo/conversion/converter_utils.py

Line 209 in ecdc040

get_axes_for_reduce_op = functools.partial(

gs-olive · 2023-09-25T22:01:50Z

py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py

+    if weight is None:
+        weight = np.array(1.0)
+
+    if bias is None:
+        bias = np.array(0.0)


I see - in that case, it might be preferable to use to_numpy(0.0), for instance, to get back a default-formatted numpy array for the float default. Additionally, I noticed the code below has some issues:

gamma = to_numpy(weight.reshape(*shape)) ##### Above is invalid, since the reshape should apply to the numpy output. It should instead be: gamma = to_numpy(weight).reshape(shape)

The same as the above applies for beta.

Additionally, lines 194 - 196 should be using get_axes_for_reduce_op, as here:

TensorRT/py/torch_tensorrt/dynamo/conversion/converter_utils.py

Line 209 in ecdc040

get_axes_for_reduce_op = functools.partial(

gs-olive · 2023-09-25T22:02:47Z

py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py

+    weight: Optional[Union[TRTTensor, torch.Tensor, np.ndarray]],
+    bias: Optional[Union[TRTTensor, torch.Tensor, np.ndarray]],


TRTTensor would not be a valid input here, for the scale layer

Do you mean the type of weight and bias in all the three functions should be Optional[Union[torch.Tensor, np.ndarray]]? I see its native function:

func: layer_norm(Tensor input, SymInt[] normalized_shape, Tensor? weight=None, Tensor? bias=None, float eps=1e-05, bool cudnn_enable=True) -> Tensor

Yes, I think it should be Optional[Union[torch.Tensor, np.ndarray]], because if either of those is a TRTTensor, the computation below would not work (to_numpy can't be called on a TRTTensor)

gs-olive · 2023-09-27T21:04:46Z

As discussed, add group_norm implementation here. Additionally, for any converters added, remove those converters from torch_tensorrt.dynamo.lowering._decomposition_groups

zewenli98 · 2023-09-28T23:03:52Z

@gs-olive group_norm was added!

gs-olive

Added a few comments. Additionally, if the dynamic shape version of this converter is not passing, that is okay since it is not required for the first pass of support

py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py

gs-olive · 2023-09-28T23:48:04Z

py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py

    scale = cast(torch.Tensor, to_numpy(weight)) / np.sqrt(
-        cast(torch.Tensor, to_numpy(running_var)) + cast(float, eps)
+        cast(torch.Tensor, to_numpy(running_var)) + eps


The torch.Tensor cast can be removed, because to_numpy will return an np.ndarray, so this typing would be incorrect.

gs-olive · 2023-09-28T23:51:09Z

py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py

+    eps_field = trt.PluginField(
+        "eps", np.array(eps, dtype=np.float32), trt.PluginFieldType.FLOAT32
+    )
+    num_groups_filed = trt.PluginField(
+        "num_groups", np.array(num_groups), trt.PluginFieldType.INT32
+    )
+
+    field_collection = trt.PluginFieldCollection([eps_field, num_groups_filed])
+
+    try:
+        # Here's the schema of the plugin:
+        # https://github.com/NVIDIA/TensorRT/blob/release/8.6/plugin/groupNormalizationPlugin/GroupNormalizationPlugin_PluginConfig.yaml
+        plugin = get_trt_plugin("GroupNormalizationPlugin", field_collection, "1")
+    except AssertionError:
+        _LOGGER.error(
+            "Unable to find group norm plugin, fall back to TensorRT implementation."
+        )
+
+    layer = network.add_plugin_v2([input, scale, bias], plugin)
+    set_layer_name(layer, target, f"{name}_GroupNormalizationPlugin", source_ir)
+
+    # PyTorch requires three return values: (out, mean, rstd)
+    dummy_tensor = torch.tensor(0)
+    return layer.get_output(0), dummy_tensor, dummy_tensor


Is it possible to avoid invoking the plugin here, and instead use the full implementation, adapting from here: https://github.com/NVIDIA-AI-IOT/torch2trt/blob/36656b614f3fbc067ac673932e2200d7afdae712/torch2trt/converters/group_norm.py#L7-L73? The plugin is not preferable for use in new converters unless it cannot be otherwise supported.

Alternatively, the TRT layer-based implementation can be the backup for the plugin, etc.

gs-olive · 2023-09-28T23:51:42Z

py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py

+    eps_field = trt.PluginField(
+        "eps", np.array(eps, dtype=np.float32), trt.PluginFieldType.FLOAT32
+    )
+    num_groups_filed = trt.PluginField(
+        "num_groups", np.array(num_groups), trt.PluginFieldType.INT32
+    )
+
+    field_collection = trt.PluginFieldCollection([eps_field, num_groups_filed])
+
+    try:
+        # Here's the schema of the plugin:
+        # https://github.com/NVIDIA/TensorRT/blob/release/8.6/plugin/groupNormalizationPlugin/GroupNormalizationPlugin_PluginConfig.yaml
+        plugin = get_trt_plugin("GroupNormalizationPlugin", field_collection, "1")
+    except AssertionError:
+        _LOGGER.error(
+            "Unable to find group norm plugin, fall back to TensorRT implementation."
+        )
+
+    layer = network.add_plugin_v2([input, scale, bias], plugin)
+    set_layer_name(layer, target, f"{name}_GroupNormalizationPlugin", source_ir)
+
+    # PyTorch requires three return values: (out, mean, rstd)
+    dummy_tensor = torch.tensor(0)
+    return layer.get_output(0), dummy_tensor, dummy_tensor


The returned values here should be correct intermediate tensors from during the computation unless we explicitly remove support for nodes which need the other two values

gs-olive · 2023-09-29T00:12:24Z

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.native_layer_norm.default)  # type: ignore[misc]


Based on the schema of native_layer_norm, it looks like it requires 3 outputs much like native_group_norm. As a comment on both of those - if you want to support it with essentially the same converter as the regular layer norm, you can do the following:

Add this validator

def validator(layer_norm: Node) -> bool: # Validate only one user, which is a getitem node that accesses the first element in the list return (len(layer_norm.users) == 1 and list(node.users)[0].target == operator.getitem and list(node.users)[0].args[1] == 0))

Add this converter

@dynamo_tensorrt_converter(torch.ops.aten.native_layer_norm.default, capability_validator=validator) def converter(...): return (regular_layer_norm, )

It is important that the above converter returns a tuple, because it will be accessed by getitem, but as you have validated, it will only access the first element. This should also work for group norm.

zewenli98 · 2023-09-30T02:48:59Z

This PR depends on #2347, #2354 and #2355

gs-olive · 2023-10-03T00:05:25Z

@zewenli98 - when you have the chance, please rebase this PR to the latest main. Additionally, to follow up on the discussion from this comment, the individual functions for layer_norm, batch_norm, etc. should likely return their relevant intermediate values too, so that we can convert those native_layer_norm-style functions.

zewenli98 · 2023-10-03T00:09:14Z

Yes! It's still in progress. Thanks for the reminder!

support group norm, and improve batch and layer norms

gs-olive

Looks good to me - will update again pending a manual check against SD

gs-olive

Works on SD - looks good to me!

facebook-github-bot added the cla signed label Sep 20, 2023

github-actions bot requested a review from peri044 September 20, 2023 23:21

zewenli98 self-assigned this Sep 22, 2023

gs-olive self-requested a review September 22, 2023 18:47

gs-olive reviewed Sep 22, 2023

View reviewed changes

gs-olive reviewed Sep 23, 2023

View reviewed changes

gs-olive reviewed Sep 25, 2023

View reviewed changes

zewenli98 force-pushed the norm_dynamo_converters branch from fd820e6 to 4aa4dce Compare September 28, 2023 23:50

gs-olive requested changes Sep 28, 2023

View reviewed changes

github-actions bot added the component: lowering Issues re: The lowering / preprocessing passes label Sep 28, 2023

gs-olive reviewed Sep 29, 2023

View reviewed changes

gs-olive added the priority: high label Sep 29, 2023

zewenli98 changed the title ~~fix: update batch_norm and layer_norm~~ feat: support group_norm, batch_norm, and layer_norm Sep 29, 2023

zewenli98 added the WIP Work is in progress, pull request should not be merged yet label Oct 3, 2023

github-actions bot requested a review from apbose October 3, 2023 00:10

zewenli98 force-pushed the norm_dynamo_converters branch from 8a41cf9 to 4f585d8 Compare October 7, 2023 02:31

zewenli98 added 4 commits October 9, 2023 15:57

update batch_norm and layer_norm

c9703a7

fix bugs

c977b21

fix type bug

0b430f0

support group norm, and improve batch and layer norms

update decomposition_groups

928d172

zewenli98 added 3 commits October 9, 2023 15:57

update group_norm with native ops

d8a7c2d

rebase and update three norms

4cc353d

add decorators

84b58dd

zewenli98 force-pushed the norm_dynamo_converters branch from f628c0c to 84b58dd Compare October 9, 2023 22:57

gs-olive reviewed Oct 10, 2023

View reviewed changes

gs-olive approved these changes Oct 10, 2023

View reviewed changes

gs-olive merged commit 1c24432 into pytorch:main Oct 10, 2023
17 of 19 checks passed

gs-olive pushed a commit that referenced this pull request Oct 10, 2023

feat: support group_norm, batch_norm, and layer_norm (#2330)

738771a

This was referenced Oct 10, 2023

cherry-pick: Key converters and documentation to release/2.1 #2387

Merged

Add BatchNorm Converter Fix / Replacement for new Partitioner #2120

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: support group_norm, batch_norm, and layer_norm #2330

feat: support group_norm, batch_norm, and layer_norm #2330

zewenli98 commented Sep 20, 2023

gs-olive left a comment

gs-olive Sep 23, 2023

gs-olive Sep 23, 2023

gs-olive Sep 23, 2023

zewenli98 Sep 25, 2023

gs-olive Sep 25, 2023

gs-olive Sep 25, 2023

gs-olive Sep 25, 2023

zewenli98 Sep 28, 2023

gs-olive Sep 28, 2023 •

edited

Loading

gs-olive commented Sep 27, 2023

zewenli98 commented Sep 28, 2023

gs-olive left a comment

gs-olive Sep 28, 2023

gs-olive Sep 28, 2023

gs-olive Sep 28, 2023

gs-olive Sep 28, 2023

gs-olive Sep 29, 2023

zewenli98 commented Sep 30, 2023 •

edited

Loading

gs-olive commented Oct 3, 2023 •

edited

Loading

zewenli98 commented Oct 3, 2023

gs-olive left a comment

gs-olive left a comment

		weight: Optional[Union[TRTTensor, torch.Tensor, np.ndarray]],
		bias: Optional[Union[TRTTensor, torch.Tensor, np.ndarray]],

		)


		@dynamo_tensorrt_converter(torch.ops.aten.native_layer_norm.default) # type: ignore[misc]

feat: support group_norm, batch_norm, and layer_norm #2330

feat: support group_norm, batch_norm, and layer_norm #2330

Conversation

zewenli98 commented Sep 20, 2023

Description

Type of change

Checklist:

gs-olive left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gs-olive Sep 28, 2023 • edited Loading

Choose a reason for hiding this comment

gs-olive commented Sep 27, 2023

zewenli98 commented Sep 28, 2023

gs-olive left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zewenli98 commented Sep 30, 2023 • edited Loading

gs-olive commented Oct 3, 2023 • edited Loading

zewenli98 commented Oct 3, 2023

gs-olive left a comment

Choose a reason for hiding this comment

gs-olive left a comment

Choose a reason for hiding this comment

gs-olive Sep 28, 2023 •

edited

Loading

zewenli98 commented Sep 30, 2023 •

edited

Loading

gs-olive commented Oct 3, 2023 •

edited

Loading