Norms Refractor

[phu0ngng]

Description

Layernorm and RMS refractor, preparing for TE/cuDNN Norms Integration.

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

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

[phu0ngng]

/te-ci pytorch

[timmoon10]

Suggested change

	#ifndef TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_H_
	#define TRANSFORMER_ENGINE_COMMON_LAYER_NORM_LN_H_
	#ifndef TRANSFORMER_ENGINE_COMMON_LAYER_NORM_NORMS_H_
	#define TRANSFORMER_ENGINE_COMMON_LAYER_NORM_NORMS_H_

[timmoon10]

I think it would be cleaner if we wrapped these objects in a namespace:

Suggested change

	namespace transformer_engine {
	namespace transformer_engine {
	namespace norms {

The names of these classes (LaunchParams, FwdParams, etc) are not very specific, so putting them in a namespace helps make clear that they are related to LayerNorm/RMSNorm.

[timmoon10]

This zeroing might be unnecessary. The FP8 scale update kernel called by PyTorch also zeros out the amax.

That said, this matches the existing behavior so no need to change right now.

[timmoon10]

Are we using this value?

Suggested change

int scale_byte_size;

[timmoon10]

If we remove the amax zeroing kernel, I don't think we are using this value:

Suggested change

int amax_byte_size;

If we want to keep an option to zero out the amax, then could we just assume the amax is FP32 like we do in the kernels?

[timmoon10]

Logically this can be split into three enums:

Suggested change

	enum NVTE_NORM_TYPE {
	LN_FWD_TE,
	LN_BWD_TE,
	LN_FWD_CUDNN,
	LN_BWD_CUDNN,
	RMS_FWD_TE,
	RMS_BWD_TE,
	RMS_FWD_CUDNN,
	RMS_BWD_CUDNN,
	};
	enum class NormType { LayerNorm, RMSNorm };
	enum class NormStage { Forward, Backward };
	enum class NormImpl { TE, CUDNN };

[timmoon10]

Suggested change

	/* #include <transformer_engine/layer_norm.h> */

[timmoon10]

While templating on IF_TUNED makes sense for our current kernel implementation, I don't think it's a good idea going forward. The reason TE compilation is so slow (and is unusable on some systems) is because we statically compile so many tuned LayerNorm/RMSNorm kernels. If we choose to keep the TE tuned kernels, we should port them to NVRTC so we can compile them at runtime, similar to how we handle the transpose kernels. This would completely remove the need for this static registry for tuned kernels (we'll still need it for general kernels though).

[timmoon10]

I see that calling this constructor is considered an error:

TransformerEngine/transformer_engine/common/layer_norm/norms.cpp

Lines 147 to 152 in 41273eb

	template <NVTE_NORM_TYPE NormEnum>
	NormFwdTe<NormEnum>::NormFwdTe() {
	if constexpr (NormEnum == NVTE_NORM_TYPE::LN_FWD_TE) {
	NVTE_ERROR("NormFwdTe default constructor is only for its inherited classes!");
	}
	}

We can catch this at compile-time with:

Suggested change

	NormFwdTe();
	NormFwdTe() = delete;

[timmoon10]

We can catch these errors at compile-time with a static_assert:

Suggested change

	if constexpr (!IF_TE_FWD_NORMS<NormEnum>()) NVTE_ERROR("Unexpected NVTE_NORM_TYPE!");
	static_assert(IF_TE_FWD_NORMS<NormEnum>(), "Unexpected NVTE_NORM_TYPE");

I see this if constexpr (...) { NVTE_ERROR(...) } pattern used in several other places.