How to avoiding having too large checkpoints in CLU?

[marcvanzee]

Original question by @linlu-qiu:

Hi, CLU question: i'm using MultihostCheckpoint and noticed the size of my checkpoint (ckpt-*.data) keeps increasing until it reaches some maximum value (which seems to be related to the checkpoint every interval I define, e.g. checkpoint every 200 steps leads to larger size than checkpoint every 100 steps). This causes memory exceed error when I increase the checkpoint interval. I tried to use standard flax.training.checkpoints and it works fine. Any idea what might be wrong?

[marcvanzee]

Answer by @Marvin182:

the ckpt-*.data files contain the state of your input pipeline. This is important if you care about deterministic training. If not you can stop checkpointing the input iterator or switch to flax.training.checkpoints (which will not checkpoint the input pipeline).