[PyTorch] Recipe heuristics for initializing quantized weights

[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 a heuristic option to the quantization recipes, with support for "performance" and "inference". In the future, we may want to consider a "memory" heuristic that prioritizes memory usage over performance.

This is similar to the heuristic API in #1300.

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

• Add recipe heuristics for initializing quantized weights

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

[timmoon10]

Destroying unnecessary usages was causing problems when alternating between training steps (column-wise data needed) and validation steps (column-wise data not needed). See #1832 (comment).

[ptrendx]

TBH this issue is just because of optimizer not doing the right job with quantizing. If we made it so it uses the quantizer then we would not need this part at all.

[timmoon10]

The layers will configure the quantizer to avoid unnecessary allocations:

TransformerEngine/transformer_engine/pytorch/module/linear.py

Lines 225 to 233 in 855fa65

	# Configure quantizer
	if weight_quantizer is not None:
	columnwise_usage = is_grad_enabled and inp.requires_grad
	if not columnwise_usage:
	columnwise_usage = (
	is_fp8_activation_recompute_enabled()
	and not in_fp8_activation_recompute_phase()
	)
	weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage)

This is what we want when allocating new buffers, but is overly aggressive when dealing with an existing QuantizedTensor. We could remove this logic from get_weight_workspace, but I don't like how it would ignore the configuration within the quantizer.

[timmoon10]

/te-ci pytorch

[ptrendx]

The name is not the best - wouldn't you want performance during inference?

[timmoon10]

Not necessarily if you're memory constrained.

Perhaps a naming scheme like "training_performance", "inference_performance", "training_memory", "inference_memory" would be more precise?

[ptrendx]

This makes me think that internal should maybe be an option to tex.quantize rather than the member of quantizer.

[timmoon10]

I have mixed opinions.

• The layers know which tensors can be internal tensors and which must be PyTorch tensors. internal seems like a usage hint just like whether it needs row-wise/column-wise data.
• We override internal multiple times, enough to make it feel redundant. These are usually in special cases outside a layer's normal operation (when primary weights are quantized, when setting tensor.data).

Maybe tex.quantize should have an option to force internal=False, but otherwise respect the quantizer's config? This seems a little overcomplicated though.

[timmoon10]

#1847 handles the main issue this PR was addressing: unnecessary memory usage when initializing quantized weights for use in inference. This API is more general though and we may revisit in the future for specialized use-cases, e.g. when it is worth sacrificing performance for reduced memory usage.