[PyTorch] Inference mode disables initializing quantized weights with column-wise usage

[timmoon10]

Description

When initializing a model with quantized weights, the required data is different for training and inference (training requires row-wise data for forward GEMM and column-wise data for dgrad GEMM, inference only requires column-wise). This PR adds logic so the model will only initialize quantized weights with the data required for inference when initialized within no-grad mode or inference mode. It is also less aggressive about deallocating weight data in order to handle cases where we were alternating between training and validation modes.

This is an alternative to #1827. The heuristic API in that PR are somewhat redundant with these plain PyTorch APIs, but it is also more general.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Avoid initializing quantized weights with column-wise usage if grads are not enabled
• Do not deallocate unnecessary usages in weight tensors during forward pass

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[timmoon10]

/te-ci pytorch

[ksivaman]

Why is this needed?

[timmoon10]

I find it convenient to be able to access a class without explicitly doing from ... import ...:

TransformerEngine/tests/pytorch/test_sanity.py

Line 1393 in 649b04c

module = transformer_engine.pytorch.ops.Linear(hidden_size, hidden_size)

It's just a matter of style though. Within the package we explicitly list the imports to order to guarantee only relative imports, but this isn't relevant for tests since we always do absolute imports. Also, Google's style guide recommends against it.

[timmoon10]

/te-ci pytorch

[ksivaman]

LGTM