Delete extra tensor objects after restoring float8 tensors

[sudhakarsingh27]

Description

After restoring the float8 tensors in the backward passes of LayernormMLP, LayernormLinear and Linear, the tensor objects are not needed but ctx.tensor_objects still holds the reference and hence it results in extra memory usage. This fixes it.

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:

• Delete extra reference to tensor_objects once they're used in the backwards of LayernormMLP, LayernormLinear and Linear modules.

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

[ksivaman]

@pggPL IIRC you removed these during a numerics debugging effort, do you remember why?

[pggPL]

If weight is in fp8 I want to have it in save_for_backward() - for offloading. If there is forward, but backward is not invoked, it will result in removing the weight. I discussed it with @ptrendx and he proposed some solution with flag internal in prepare_for_saving - to set it True if tensor is not owned and remove tensors iff they are internal. It seems that we forgot about this.

[ptrendx]

Ok, so that would be solved by overriding this function in Float8Tensor and MXFP8Tensor to just return self and None instead.

Also, in https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/tensor/quantized_tensor.py#L30 why do we check for exactly Tensor or Param and not just isinstance(torch.Tensor)? This should solve this as well, right?

[pggPL]

Now logic of restoring tensor is inside the tensor object. If tensor object is None, we assume that this was standard torch.tensor. If it is QuantizedTensor, then it object is responsible for restoring itself, so we need to somehow save it.

[ptrendx]

Well, QuantizedTensor is in a way a standard tensor - at least it can be passed whole through save_for_backward, so there is nothing to restore afterwards.

[pggPL]

ok, it makes sense

[sudhakarsingh27]

/te-ci