Handle the scaling factor when amax is too tiny that leads to an infinite scale

[jinzex]

When the amax is too tiny that the scale becoming infinite in FP32, we set the scale to the max value of FP32. In this case, the tensor’s amax won't get mapped to the FP8 max representable, but rather something below that, but this is the best thing we can do.

cc @Oleg-Goncharov

[jinzex]

@timmoon10 Thanks Tim for your review and suggestions. All review suggestions were applied.

[jinzex]

@ksivaman Currently test_recipe.py is not included in the qa/L0_pytorch_unittest/test.sh. And test_amax_and_scale_update is failed for the case that has is_first_microbatch=False, could you give some suggestions on how to fix it?

[ksivaman]

@jinzex Could you add the test_recipe.py in the qa folder such that it is included in our tests? We should also remove freezing scaling factors during the is_first_microbatch tests since it was removed in #575. I suspect this is also the reason for the failures you observe.

[jinzex]

@ksivaman Thanks for your reply! Do you mean removing the is_first_microbatch=False case?

TransformerEngine/tests/pytorch/test_recipe.py

Lines 31 to 39 in f85553e

	@pytest.mark.parametrize("amax_history_len", [1, 31, 1024])
	@pytest.mark.parametrize("amax_compute_algo", ["max", "most_recent"])
	@pytest.mark.parametrize("is_first_microbatch", [None, True, False])
	def test_amax_and_scale_update(
	self,
	amax_history_len: int,
	amax_compute_algo: str,
	is_first_microbatch: Optional[bool],
	margin: int = 2,

For example, change it to

    @pytest.mark.parametrize("amax_history_len", [1, 31, 1024])
    @pytest.mark.parametrize("amax_compute_algo", ["max", "most_recent"])
    @pytest.mark.parametrize("is_first_microbatch", [None, True])
    def test_amax_and_scale_update(
        self,
        amax_history_len: int,
        amax_compute_algo: str,
        is_first_microbatch: Optional[bool],
        margin: int = 2,
    ):

[ksivaman]

I meant that in this line, we should set update_weight_scale_inv to True and check the results!

[jinzex]

Thanks for the suggestion! That fixed the test. Changes has been committed.

[timmoon10]

I don't think we should change test_recipe.py. The test failure exposes a valid bug introduced by the recipe change in #575. In particular, by updating the weight scales in every microbatch step, we might change the FP8 scale in a step where we do not change the FP8 data.

This bug is beyond the scope of this PR though, and otherwise this PR looks good. I think a better solution is to merge this PR without making any changes to the tests and we will fix that bug in an upcoming PR.

[timmoon10]

/te-ci pytorch

[jinzex]

@timmoon10 Thanks Tim for the comment! The changes to update_weight_scale_inv have been reverted.

[timmoon10]

LGTM. This will expose some test failures from an existing bug (see #786 (review)) and I see some unrelated linter failures (see #816).

[timmoon10]

/te-ci pytorch