Unexpectedly high CPU memory usage when running GSM8K with LLaDA

[XiangZhang-zx]

Hi,
I noticed very high CPU memory usage when running GSM8K with the official scripts.

• Single process (1 GPU): ~160GB RAM
• 4 GPUs: ~600GB RAM
  Is this expected behavior, or might it indicate a memory leak in the dataloader/diffusion loop?

Any advice would be greatly appreciated!

[jinxinzhou]

I meet same OOD issue. After examing the code, I think this issue come from the map() function at line 256-257 in eval_llada.py, as the tokenized prompt will cost lots of memory.

I remove the lines 256-257

        ds = ds.map(_tokenize)
        ds = ds.with_format("torch")

and change the line 261 to

            elem = _tokenize(elem)
            prompt = torch.tensor(elem["question"], device=self.device).unsqueeze(0)

[nieshenx]

Thanks for your interest in our work.

You can refer to this PR for details.