Remove CUDA Sync Workaround in Evo2 Training Script

[gagank1]

There is a bug in TransformerEngine's LayerNormLinear where self.layer_norm_weight may not be fully initialized before the first forward pass due to asynchronous CUDA execution. It can contain garbage values instead of the expected initial values, causing NaNs and other potential issues. We have a workaround in place that sets CUDA_LAUNCH_BLOCKING=1 before the first batch and removes it after to force synchronization on the first batch only. This needs to be removed once the bug is fixed upstream.

[jstjohn]

@gagank1 interestingly I'm having the same issue with a non-TE layer now:

bionemo-framework/bionemo-recipes/recipes/evo2_megatron/src/bionemo/evo2/models/megatron/hyena/hyena_utils.py

Lines 1247 to 1253 in 9dbcd2b

	def reset_parameters(self):
	"""Reset the parameters of the ParallelShortHyenaOperator."""
	with get_cuda_rng_tracker().fork(), torch.no_grad():
	bounds = math.sqrt(1 / self.kernel_size)
	conv_init_method = partial(torch.nn.init.uniform_, a=-bounds, b=bounds)
	self.conv_bias.data = conv_init_method(self.conv_bias.data)
	self.short_conv.reset_parameters()

Even though I reset the layer 4x during init currently, it only actually works if either:

1. I put pdb statements after each call to reset_parameters() which seems to block long enough
2. I run the job with CUDA_LAUNCH_BLOCKING like you mentioned.

So it's not just megatron/TE, wonder what's going on?

[jstjohn]

This is now fixed in megatron with NVIDIA/Megatron-LM#2652

[jstjohn]

Addressed in megatron bridge here: NVIDIA-NeMo/Megatron-Bridge#1706