Fixed memory issues during forward

[AleHD]

Nanotron seems to consume disproportionately more memory on its activations compared to megatron. This is due to at least the following factors:

• The glu activation, which is not fused, allocate two tensors: the activation and the element-wise multiplication. Fusing this operation provides (relatively small) memory gains.
• During the differentiable operations (specifically, the DifferentiableAllGather and DifferentiableReduceScatterSum) a new tensor is allocated via torch.empty. This tensor is cached through the entire forward pass until the backward pass. However, this cache is unnecessary as these tensors are not used at all (rather, they are reconstructed via communication during the backward). Getting rid of these allocations provide significant memory gains. To fix this, this PR introduces a global memory buffer (MemoryBuffer singleton) that recycles the allocated spaces, similar to megatron.

Attached: Memory traces of the default nanotron implementation (which OOMs), the current PR implementation and megatron. The memory traces represent the first rank of a tp8 pp4 dp1 llama70b.

[AleHD]

I reviewed hf's comments and made appropiate adjustments. See: huggingface#203 (comment). Should be ready for review :).

[ischlag]

Tried to run it but I get

AttributeError: 'dict' object has no attribute '_is_using_mup' self.model_config._is_using_mup = isinstance(self.init_method, SpectralMupInit)self.model_config._is_using_mup = isinstance(self.init_method, SpectralMupInit)
Note, I used the 70B config from the pretrain repo but I don't see any meaningful differences to the config. And there are a lot of other changes in the PR seemingly unrelated to the topic..?

[AleHD]

Might have to do with the version of the main branch. I will update everything on my end and make sure that it works properly before opening the PR again. And yes, there are some seemingly unrelated changes. My fork is from upstream nanotron so those are commits pushed to the main branch upstream. We should update our fork to make it easier to review.