Async amax reduction

[erhoo82]

• Support async AMAX reduction
• AMAX reduction is default enabled and can be disabled by setting an env var (NVTE_ASYNC_AMAX_REDUCTION)

[ptrendx]

@erhoo82 Please sign your commits (please refer to CONTRIBUTING.rst to see how to do this).

[erhoo82]

Thanks, signed the commit.

[timmoon10]

/te-ci

[timmoon10]

/te-ci

[erhoo82]

@ksivaman
I tested and confirmed the functionality.

[ksivaman]

LGTM, thanks!

[ksivaman]

/te-ci

[ksivaman]

/te-ci

[ksivaman]

/te-ci

[Victarry]

Hello, @ksivaman @erhoo82 , I think it's more reasonable to put amax_reduce_handle_fwd.wait() before FP8GlobalStateManager.copy_amax_from_global_buffer(self.fp8_meta, forward=True), which ensures the reduction is finished before copy amax from global buffer.

TransformerEngine/transformer_engine/pytorch/module/base.py

Lines 582 to 616 in b8ba734

	# Previous iteration was grad_enabled
	if self.fp8_meta.get("update_amax_and_scale_fwd", False):
	if self.fp8_meta["recipe"].reduce_amax:
	FP8GlobalStateManager.copy_amax_from_global_buffer(self.fp8_meta, forward=True)
	amax_and_scale_update(
	self.fp8_meta, True, update_weight_scale_inv=update_weight_scale_inv
	)
	FP8GlobalStateManager.set_amax_buffer_key_deletion(self.fp8_meta, forward=True)
	else:
	amax_and_scale_update(
	self.fp8_meta, True, update_weight_scale_inv=update_weight_scale_inv
	)
	if self.fp8 and self.training:
	# Setup for amax reduction
	if self.fp8_meta["recipe"].reduce_amax:
	self.fp8_meta["first_module"] = FP8GlobalStateManager.is_first_fp8_module()
	if self.fp8_meta["first_module"]:
	# Wait for the prior AMAX reduction to finish
	amax_reduce_handle_fwd = FP8GlobalStateManager.get_amax_reduce_handle_fwd()
	if amax_reduce_handle_fwd is not None:
	amax_reduce_handle_fwd.wait()
	self.fp8_meta["autocast_id_fwd"] = (
	FP8GlobalStateManager.new_fp8_context_id())
	FP8GlobalStateManager.set_fp8_context_id(self.fp8_meta["autocast_id_fwd"])
	else:
	self.fp8_meta["autocast_id_fwd"] = (
	FP8GlobalStateManager.get_fp8_context_id())
	self.fp8_meta["autocast_id_fwd_stack"].append(
	self.fp8_meta["autocast_id_fwd"]
	)
	FP8GlobalStateManager.add_amax_to_global_buffer(self.fp8_meta, forward=True)
	self.fp8_meta["update_amax_and_scale_fwd"] = True
	else:
	self.fp8_meta["update_amax_and_scale_fwd"] = False