Fix a crash in NeMo 2.0 during module._apply(lambda t: t.cpu())

[guyueh1]

Description

In Nemo 2.0 during job exit, lightning calls a module._apply(lambda t: t.cpu()) on the GPT model which triggers an illegal memory access error in the TE dequantize kernel. This PR fixes the issue.

Fixes # (issue)

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 refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

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]

Can you explain the race condition that this is fixing? From what I can tell, Float8Tensor.cpu should already synchronize the GPU:

TransformerEngine/transformer_engine/pytorch/tensor/quantized_tensor.py

Lines 321 to 323 in e70f913

	def cpu(self, memory_format=torch.preserve_format) -> torch.Tensor:
	# pylint: disable=missing-function-docstring
	return self.dequantize().cpu(memory_format=memory_format)

The actual problem is later in Float8Tensor._set_data:

TransformerEngine/transformer_engine/pytorch/tensor/float8_tensor.py

Line 528 in e70f913

self.data = self._quantizer.quantize(tensor)

We are passing a CPU tensor into the quantize kernel, and I don't think we ever move it to GPU. This doesn't explain why this PR fixes the IMA, so I could have missed something.

If my interpretation is the actual root cause, the quickest fix is to modify Float8Tensor._set_data with:

self.data = self._quantizer.quantize(tensor.to(device=self.device))

More long-term fixes are to handle CPU tensors in the quantize function or to support CPU Float8Tensors.

[guyueh1]

@timmoon10
Re: We are passing a CPU tensor into the quantize kernel, and I don't think we ever move it to GPU.
I checked the device of 'tensor' in debug mode, after torch.cuda.synchronize() it is on cpu, however, the self._quantizer.quantize(tensor) worked fine. So there should be some code to handle cpu or move it to gpu somewhere.

I also don't fully understand the race condition, it just happened to work. I will dig into it, and reply here my findings.

[guyueh1]

Findings:

1. Adding torch.cuda.synchronize() at beginning of _set_data works (the original proposal of this PR)
2. @timmoon10 's suggestion (self.data = self._quantizer.quantize(tensor.to(device=self.device))) also works
   Still digging why IMA happens

[guyueh1]

I decide to revert the torch.cuda.synchronize() change because I can't understand why it would work; I apply a new way to fix it by making sure the tensor is moved to self.device if it were on CPU. I confirmed it fixed the IMA.
@timmoon10 what you think about the current version

[timmoon10]

/te-ci pytorch

[timmoon10]

LGTM