Changes to make torch.compile work

[zalbanob]

Key Performance Metrics Analysis:

Memory Distribution:

• System RAM maintains ~400GB utilization
  

• GPU memory utilization is notably imbalanced:

  • GPUs 0-2: ~60%
  • GPU 3: >90% (hosting two encoder-decoder transformers)
    

GPU Computational Load:

• Utilization metrics:
  • GPUs 0-2: ~12% mean utilization
  • GPU 3: 25% mean utilization

Training Throughput:

• Steady-state performance: 58.48 samples/second (mean)
• Performance distribution shows strong consistency (IQR: 3.8)
• 90% of samples fall between 56.7-71.7 samples/second
  
  

Current Bottleneck: The assignment of two encoder-decoder transformers on GPU 3 is creating a memory constraint that prevents batch size scaling, and subsequently limiting our ability to increase the number data loaders.