[PyTorch] Normalization ops

[timmoon10]

Description

This PR extends the operation-based API (see #707) with LayerNorm, RMSNorm, and FP8 cast operations.

Compare with the existing module-based API:

# Module-based API
module1 = te.LayerNormLinear(...)

# Operation-based API
module2 = te.ops.Sequential(
    te.ops.LayerNorm(...),
    te.ops.Linear(...),
)

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:

• LayerNorm operation
• FP8 cast operation
• RMSNorm operation

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]

/te-ci pytorch

[timmoon10]

/te-ci pytorch

[timmoon10]

/te-ci pytorch

[timmoon10]

/te-ci pytorch

[timmoon10]

/te-ci pytorch

[timmoon10]

/te-ci pytorch

[ptrendx]

This seems a little backwards. if dtype is supposed to be the new argument name, then why is it in kwargs? Both params_dtype and dtype should be regular parameters, there should be a deprecation warning when somebody uses params_dtype and also the check for duplicate assignment like the one you have here.
Also, similar treatment should be done for hidden_size and sequence_parallel (especially the last one seems to be just gone completely so there should be some explanation that it was unused before or something?)

[timmoon10]

My thinking is that we should forward kwargs directly to te.ops.LayerNorm as much as possible so that we only have to change the API in one place if we ever make changes in the future. We include the deprecated options as explicit kwargs since they are specific to the module.

This function signature also maintains backward compatibility for users who pass in the options as positional args, e.g.:

TransformerEngine/tests/pytorch/test_onnx_export.py

Lines 676 to 678 in 0ee5ccd

	te.LayerNorm(
	inp_shape[1], eps, params_dtype=dtype, zero_centered_gamma=zero_centered_gamma
	)

[ptrendx]

So why is sequence_parallel option deprecated then? I believe Megatron is using those to guide their logic for the optimizer. I know it is not great, but we should not break them.

[timmoon10]

Maybe "legacy" is better. We should treat this as a weird, Megatron-specific integration.

[ptrendx]

Same issue as in LN.

[ptrendx]

I believe you said that it is mostly an utility op for tests, right? We should probably mention that in this documentaiton.

[ptrendx]

Also maybe we should consider generalizing it a bit with a name that is not specific to FP8? (Like just Quantize)?

[timmoon10]

It could be a helpful op for users as well. For example, if a user wants to have discrete layers for design reasons but still wants to fuse some operations with FP8 casts:

act = te.ops.Sequential(te.ops.GeLU(), te.ops.Quantize())
linear = te.ops.Sequential(te.ops.Linear())
y = act(x)
z = linear(y)

[timmoon10]

/te-ci pytorch