Replaces torch.linalg.norm with torch.linalg.vector_norm for vector-norm

`torch.linalg.norm` supports various calculations based on `dim` parameter:
- If dim is an int, the vector norm will be computed.
- If dim is a 2-tuple, the matrix norm will be computed.
- If dim=None and ord=None, A will be flattened to 1D and the 2-norm of the resulting vector will be computed.
- If dim=None and ord!=None, A must be 1D or 2D.

Therefore, vector norm is not computed when `dim` is tuple. (Nit: `torch.linalg.vector_norm` is more explicit for vector norms.)

Author: namgyu-youn

modelopt/torch/nas/plugins/megatron.py

@@ -622,7 +622,7 @@ def _estimate_all_head_importance(self) -> TracedHp.Importance:
     def _estimate_query_group_importance(self) -> TracedHp.Importance:
         """Return the importance of the ``num_query_groups`` hparam."""
         assert self._activations is not None, "No activations collected for importance estimation."
-        group_importance = torch.linalg.norm(
+        group_importance = torch.linalg.vector_norm(
             self._activations.view(
                 self.get_hparam("num_heads_per_group").max,
                 self.get_hparam("num_query_groups").max,