[BUG] 60GB of RAM but getting - RuntimeError: [enforce fail at alloc_cpu.cpp:75] err == 0. DefaultCPUAllocator: can't allocate memory: you tried to allocate ...

I am executing the following command to train HuggingFace's Dolly model and I am getting the following error which does not make sense because I have nearly 60GB free ram. The error is in the cpu_adam.py file:

Command:

deepspeed --num_gpus=1 --module training.trainer --deepspeed config/v100_config.json --local-output-dir testOutputDir --per-device-train-batch-size 2 --input-model EleutherAI/pythia-2.8b --logging-steps 10 --save-steps 200 --save-total-limit 20 --eval-steps 50 --warmup-steps 50^C-test-size 200 --lr 5e-6 --epochs 2

Error

...
Time to load utils op: 0.16718578338623047 seconds
Parameter Offload: Total persistent parameters: 1070080 in 258 params
2023-07-03 04:20:16 ERROR [__main__] main failed
Traceback (most recent call last):
  File "/home/admin/dolly/training/trainer.py", line 332, in <module>
    main()
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/click/core.py", line 1130, in __call__
    return self.main(*args, **kwargs)
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/click/core.py", line 1055, in main
    rv = self.invoke(ctx)
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/click/core.py", line 1404, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/click/core.py", line 760, in invoke
    return __callback(*args, **kwargs)
  File "/home/admin/dolly/training/trainer.py", line 324, in main
    train(**kwargs)
  File "/home/admin/dolly/training/trainer.py", line 279, in train
    trainer.train()
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/transformers/trainer.py", line 1649, in train
    ignore_keys_for_eval=ignore_keys_for_eval,
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/transformers/trainer.py", line 1760, in _inner_training_loop
    self.model, self.optimizer, self.lr_scheduler
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/accelerate/accelerator.py", line 1178, in prepare
    result = self._prepare_deepspeed(*args)
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/accelerate/accelerator.py", line 1505, in _prepare_deepspeed
    engine, optimizer, _, lr_scheduler = deepspeed.initialize(**kwargs)
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/deepspeed/__init__.py", line 135, in initialize
    config_params=config_params)
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/deepspeed/runtime/engine.py", line 340, in __init__
    self._configure_optimizer(optimizer, model_parameters)
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/deepspeed/runtime/engine.py", line 1298, in _configure_optimizer
    self.optimizer = self._configure_zero_optimizer(basic_optimizer)
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/deepspeed/runtime/engine.py", line 1626, in _configure_zero_optimizer
    communication_data_type=self.communication_data_type)
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/deepspeed/runtime/zero/stage3.py", line 312, in __init__
    self._setup_for_real_optimizer()
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/deepspeed/runtime/zero/stage3.py", line 371, in _setup_for_real_optimizer
    self.initialize_optimizer_states()
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/deepspeed/runtime/zero/stage3.py", line 938, in initialize_optimizer_states
    self._optimizer_step(i)
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/deepspeed/runtime/zero/stage3.py", line 858, in _optimizer_step
    self.optimizer.step()
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/torch/optim/optimizer.py", line 140, in wrapper
    out = func(*args, **kwargs)
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/admin/train-dolly/env/lib/python3.7/site-packages/deepspeed/ops/adam/cpu_adam.py", line 178, in step
    device=device)
RuntimeError: [enforce fail at alloc_cpu.cpp:75] err == 0. DefaultCPUAllocator: can't allocate memory: you tried to allocate 4012482560 bytes. Error code 12 (Cannot allocate memory)``

Start with the training command and config/v100_config.json, then trace the failure from training/trainer.py into deepspeed/ops/adam/cpu_adam.py. Reproduce the allocation error and inspect the surrounding DeepSpeed ZeRO Stage 3 optimizer setup. Done would require a confirmed cause and either a targeted fix or clear guidance for this configuration.