[refactor] Update Ln-norm logic for upcoming PyTorch update

[namgyu-youn]

What does this PR do?

Type of change: Refactor

Overview: Based on PyTorch - docs, torch.norm() will be removed in the foreseeable version. Therefore, this update replaces torch.norm() with torch.linalg.vector_norm()

Usage

# Old (current) version
torch.norm()

# New (upcoming) version
torch.linalg.vector_norm()

Testing

pytest and the comment

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines
• Is this change backward compatible?: Yes
• Did you write any new necessary tests?: Yes
• Did you add or update any necessary documentation?: No
• Did you update Changelog?: No

Additional Information

[namgyu-youn]

cc. @kevalmorabia97 @meenchen

[kevalmorabia97]

Thanks @namgyu-youn for contributing the fix. Can you also apply similar changes to usage of <tensor>.norm() throughout the repository?

Please also share testing details - whether you have run unit and GPU tests? (pytest tests/unit and pytest tests/gpu)

[namgyu-youn]

Thanks @namgyu-youn for contributing the fix. Can you also apply similar changes to usage of <tensor>.norm() throughout the repository?

@kevalmorabia97; Sure, and it seems there is no other Ln-norm logic in modelopt/. Let me know if I am missing something.

Please also share testing details - whether you have run unit and GPU tests? (pytest tests/unit and pytest tests/gpu)

The result of pytest is the following:

Also, local test code is the following:

import torch

# Create test tensors
torch.manual_seed(42)
X = torch.randn(10, 10)
X_hat, A_reg = torch.randn(10, 10), torch.randn(5, 5)

# Old Version (torch.norm)
print("=== torch.norm ===")
relative_error = torch.dist(X, X_hat, p=torch.inf) / torch.norm(X, p=torch.inf)
norm_result = torch.norm(A_reg, p=torch.inf)

print("Relative error:", relative_error.item())
print("A_reg infinity norm:", norm_result.item())
print()

# New Version (torch.linalg.vector_norm)
print("=== torch.linalg.vector_norm ===")
relative_error = torch.dist(X, X_hat, p=torch.inf) / torch.linalg.vector_norm(X, ord=torch.inf)
norm_result = torch.linalg.vector_norm(A_reg, ord=torch.inf)

print("Relative error:", relative_error.item())
print("A_reg infinity norm:", norm_result.item())

And the result is:

=== torch.norm ===
Relative error: 1.512144684791565
A_reg infinity norm: 1.9775909185409546

=== torch.linalg.vector_norm ===
Relative error: 1.512144684791565
A_reg infinity norm: 1.9775909185409546

[kevalmorabia97]

Please see https://github.com/search?q=repo%3ANVIDIA%2FTensorRT-Model-Optimizer+.norm%28&type=code

[namgyu-youn]

Please see https://github.com/search?q=repo%3ANVIDIA%2FTensorRT-Model-Optimizer+.norm%28&type=code

Got it, really thanks for your leading.

[kevalmorabia97]

Thanks for the changes. Let me take this in our internal repo where we run a more extensive CI test which is yet to be migrated to Github.

[kevalmorabia97]

Should we use vector_norm here as well?

[namgyu-youn]

This case computes vector norm because dim is int. Please check above comment; let me know if there is anything wrong in internal CI

[kevalmorabia97]

Lets use vector_norm explicitly to avoid any ambiguity and consistency with rest of the changes

[kevalmorabia97]

There is one more torch.norm change missing at

TensorRT-Model-Optimizer/modelopt/torch/nas/plugins/megatron.py

Lines 616 to 619 in bb630db

	attn_head_importance = self._activations.view(
	self.get_hparam("num_heads_per_group").max * self.get_hparam("num_query_groups").max,
	self.config.kv_channels,
	).norm(p=2, dim=1)

Can you please address that as well?

[namgyu-youn]

There is one more torch.norm change missing at

TensorRT-Model-Optimizer/modelopt/torch/nas/plugins/megatron.py

Lines 616 to 619 in bb630db

	attn_head_importance = self._activations.view(
	self.get_hparam("num_heads_per_group").max * self.get_hparam("num_query_groups").max,
	self.config.kv_channels,
	).norm(p=2, dim=1)

Can you please address that as well?

@kevalmorabia97 ; Thanks for your leading. Could you take a look at this PR?

[namgyu-youn]

Also, could you comment about #146 - Comment? Reviewing source codes in modelopt/torch/nas and modelopt/torch/prune was the start for a deeper (new feature) contribution.

But before the start, your (and other maintainers) advice must be valuable background for the directionality. I hope the ticket is still open.

[kevalmorabia97]

Your PR commit is merged. Thank you for contributing