Memory leak with Nanoset

[eliebak]

Seems like there is a memory leak that lead to a OOM on the cpu when using nanoset and training for a long time.

error summary:

slurmstepd: error: Detected 1 oom-kill event(s) in StepId=13466817.0. Some of your processes may have been killed by the cgroup out-of-memory handler.
srun: error: ip-26-0-168-30: task 6: Out Of Memory

Reproducing the issue:

• Full error: https://gist.github.com/eliebak/ed7e625afb6aa71dc5e87bebc19f03cf (from a similar job, didn't wait for this one to be OOM to save compute, but the trend in CPU memory is clearly similar)
• config: https://gist.github.com/eliebak/6adcbd12e5acc3252b107e89ab5e1ca8
• env: https://gist.github.com/eliebak/951655b5a26d9ace50fd307936b65e28
• Nanotron commit: 1e31cb9601bdff4db96a10b8f6b0b238163273e9
• Launching job with: https://gist.github.com/eliebak/81b5d42471a8521f8f587e9100b23cd7

Grafana log:
d