[C] Normalization Refactor + Adding CUDNN backend

[phu0ngng]

Description

TODO:

• Adapt normalization in JAX and Paddle.
• Benchmark performance of new APIs.

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 refractor

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

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]

We would have unexpected behavior if the user sets NVTE_BWD_LAYERNORM_USE_CUDNN=0:

Suggested change

	if (!std::getenv("NVTE_BWD_LAYERNORM_USE_CUDNN")) {
	if (!std::getenv("NVTE_BWD_LAYERNORM_USE_CUDNN")
	|| !std::atoi(std::getenv("NVTE_BWD_LAYERNORM_USE_CUDNN"))) {

Alternatively, we could use TE's getenv function:

TransformerEngine/transformer_engine/common/util/system.h

Line 22 in e5ffaa7

T getenv(const char *variable);

However, this is delicate since it can run into issues if the core lib and framework lib are compiled with different C++ ABIs. A solution might be to make getenv a header-only impl.

[phu0ngng]

Hi, I had not yet introduced changes in the framework part yet.
With the new change in PyTorch side in the latest commit, we no longer need this env check.

[timmoon10]

Nit: We use the NVTE_ prefix for C-style enums to avoid ambiguity, but it's not necessary for these C++-style enums since they're defined within the transformer_engine::normalization namespace:

Suggested change

	enum class NVTE_Norm_Backend { Te, Cudnn };
	enum class NVTE_Norm_Type { LayerNorm, RMSNorm };
	enum class NVTE_Norm_Stage { Forward, Backward };
	enum class NormBackend { Te, Cudnn };
	enum class NormType { LayerNorm, RMSNorm };
	enum class NormStage { Forward, Backward };

Not relevant yet, but in the future we might want to have separate stages for ForwardTrain and ForwardInfer if cuDNN exposes an optimized inference impl.

[timmoon10]

We should use TE's getenv function so that users can specify NVTE_FWD_LAYERNORM_USE_CUDNN=0:

Suggested change

	if (std::getenv("NVTE_FWD_LAYERNORM_USE_CUDNN")) {
	if (getenv<bool>("NVTE_FWD_LAYERNORM_USE_CUDNN")) {

See:

TransformerEngine/transformer_engine/common/util/system.h

Line 22 in c0a539c

T getenv(const char *variable);

Similar changes should be made in the backward function, as well as RMSNorm. We might also want to consider caching the value to avoid CPU overheads:

  static const bool use_cudnn_backend = getenv<bool>("NVTE_FWD_LAYERNORM_USE_CUDNN");
  if (use_cudnn_backend) {

Related: #1315 (comment)

[timmoon10]

It's not immediately obvious why we only check tensor dimensions when the workspace is empty. If it's cheap, it'll be easier to check tensors every time. Otherwise, it would be helpful to add a comment that we assume that this function is called twice (to query workspace and to launch kernel) and that we only check tensors the first time.

Similar changes should be made in the backward function, as well as RMSNorm.

[timmoon10]

Mismatched parenthesis:

Suggested change

	* y = \frac{x - E[x]}{\sqrt{Var[x] + \varepsilon}}\gamma) + \beta
	* y = \frac{x - E[x]}{\sqrt{Var[x] + \varepsilon}} \gamma + \beta

We could also make the LaTeX look nicer with:

 * y = \frac{x - \mathbb{E}[x]}{\sqrt{\text{Var}[x] + \varepsilon}} \gamma + \beta

Any changes here should also be made in the backward function and in RMSNorm.

[timmoon10]

We should provide more details since we don't include the formula in the main documentation.

Suggested change

	* \param[in] zero_centered_gamma If zero_centered_gamma is enabled
	* \param[in] zero_centered_gamma Multiply normalized values by @f$ \gamma+1 @f$ instead of @f$ \gamma @f$

Similar changes should be made to the backward function as well as RMSNorm.

[timmoon10]

Do we need to relax tolerances?

Suggested change

	double atol_bwd = 1e-3;
	double rtol_bwd = 1e-3;
	double atol_bwd = 1e-4;
	double rtol_bwd = 1e-4;

Ideally we would use tight tolerances like:

TransformerEngine/tests/pytorch/utils.py

Lines 73 to 84 in c0a539c

	if dtype == torch.float16:
	return dict(rtol=1e-3, atol=1e-5)
	if dtype == torch.bfloat16:
	return dict(rtol=1.6e-2, atol=1e-5)
	if dtype == torch.float32:
	return dict(rtol=1.3e-6, atol=1e-5)
	if dtype == torch.float64:
	return dict(rtol=1e-7, atol=1e-7)
	if dtype == torch.float8_e4m3fn:
	return dict(rtol=0.125, atol=0.0675) # epsilon = 0.0625
	if dtype == torch.float8_e5m2:
	return dict(rtol=0.25, atol=0.125) # epsilon = 0.152

[timmoon10]

Nit: Caching this value makes the code slightly more concise, but I think it's clearer and more general to just store the original enum:

Suggested change

	bool _is_layernorm;
	NVTE_Norm_Type _norm_type;

[timmoon10]

It's a bit awkward that NormalizationPlanBase includes logic for both the forward and backward kernels. I feel the "right" solution is to implement separate abstract base classes for forward and backward. That said, this approach does allow some code reuse with the cuDNN backend though, so I think it's fine for now. We can revisit in the future if needed, e.g. if we add a separate forward inference stage.

[timmoon10]

One problem with these envvars is it's all-or-nothing. This is especially troublesome for testing, where we may want to test TE and cuDNN norm kernels within the same process. Perhaps we could have some function like nvte_enable_cudnn_layernorm that sets a global variable, and the initial value is based on an envvar.

[timmoon10]

It would be helpful to test both cuDNN and TE kernels. We can reduce the number of test shapes to keep the total number of tests manageable.