Hi all, I'm trying to run the Dolly example on Metaflow but keep running into an OOM error:

OutOfMemoryError: Container killed due to memory usage

Does Metaflow enforce a hard memory limit for ECS containers? I'm using a

p3.16xlarge

and this is my resource config:

N_CPU = 32
N_GPU = 8
MEMORY = 250000

it is possible that 250GB is not enough RAM for the Dolly example. We had 400GB on our instance 🤔

Oh interesting. I thought it was only 128GB. But I'll try with 400GB

ah, good catch. It might be incorrect in that snippet. Try it with

memory=400000

. We can fix it accordingly

Okay thanks Ville. Trying it again.

that seems like an issue with the image you use and the CUDA libraries it contains

Making progress now. Just need to figure out the cuda out of memory errors:

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 256.00 MiB (GPU 3; 15.78 GiB total capacity; 14.85 GiB already allocated; 235.75 MiB free; 14.85 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

I might be getting this error because I set

bf16

to false though. But interestingly, if I set to true, it complains that volta-100 is not compatible with

bf16

This GPU issues memory issue a bit tricky and dependent on the GPU memory capacity and number of GPUs on the node interacting with deepspeed. Probably the most fruitful place to look is a light read of the deepspeed config options in the ds_config.json. At the moment, I'm not sure of how to avoid a rather slow parameter tuning loop playing around how deepspeed shares model state between GPU memory and the RAM, specifically playing with the options for offloading model state from GPU mem to RAM.

VERY helpful!!! Thanks so much Eddie!

of course!

sigh! No luck. Even a batch size of 1 results in

torch.cuda.OutOfMemoryError

you can try with

--num_gpus=3

- currently you have 16 or 32GB VRAM / GPU which might not be enough we ran our tests on an A100 that had 80GB VRAM / GPU

Ok - trying with 3 GPUs. Sadly I don't have access to an A100 (they're unavailable) and I can't even get my hands on a

p3dn.24xlarge

😭

Ahh yes, we did two tests that were working steadily • the 3 A100s on Coreweave • the 8 V100s on AWS, but these were

p3dn.24xlarge

, which can be a crapshoot to access right now so sad

note we were able to do inference with

p3.8xlarge

p3.16xlarge

. I think there are also some deepspeed options that might help here, although I am not an expert in the framework. In

ds_config,json

, there is this zero optimizer section, where you can tell deepspeed to move parameters (zero stage 3 only) and optimizer state to CPU or NVMe memory like this:

"zero_optimization": {
  "stage": 3,  
  ...

  "offload_param": {
    "device": "[cpu|nvme]",
    "nvme_path": "/local_nvme",
    "pin_memory": [true|false],
    "buffer_count": 5,
    "buffer_size": 1e8,
    "max_in_cpu": 1e9
  },
  "offload_optimizer": {
    "device": "[cpu|nvme]",
    "nvme_path": "/local_nvme",
    "pin_memory": [true|false],
    "buffer_count": 4,
    "fast_init": false
  },
  ...
}

You did inference but not training with p3.16xlarge?

yeah.

p3.16xlarge

could hold the model to do predictions, but the gradient syncing required for training was where memory issue came into play

Got it, thanks Eddie. I'll try one more time with the configuration above.

Nice, good luck 🤞

Hi Eddie, Managed to get my hands on a p3dn.24xlarge. Was surprised that it was still failing at epoch 0.14. Always seems to fail around there with Cuda OOM error. Tried batch size of 3. Could try batch size of 2 but wouldn't be very efficient and would certainly be costly as well.

Ahh shoot, sorry to hear, these GPU workflows can be frustrating! What does your deepspeed config file look like right now?

Haha no worries. I am sure my general incompetence is partly to blame. My first time training an LLM and using deepspeed. My setup is quite similar to what was recommended in the source Databricks Dolly README. Will share the config shortly. I should also note that I reduced the model to the one with only 3B parameters and I'm using the PyTorch cuda dockerfile if that matters

Ok, one low-hanging fruit we can try is to offload parameter state as well, adding another section under

"zero_optimization"

like:

"offload_param": {
     "device": "cpu",
     "pin_memory": true
}

Some other parameters we can tune to reduce max GPU memory requirements:

stage3_max_live_parameters
stage3_max_reuse_distance

I tried the first suggestion above because you raised it as a suggestion before but can certainly try again. I think the problem could be very long inputs, for that particular batch. Any idea on how to filter these long inputs out?

This line in the consts updates a max_length parameter passed to the tokenizer in trainer.py. That can control sequence length (e.g., truncate or add padding) depending on that params value. Here is a guide on how that all works generally, in our case action is going to be in trainer.py.

I could try decreasing the

max_length

a tad

Yeah, thats a decent bet. I didn't play around with that much, curious to see if that works as a quick fix just to get the thing running smoothly for you.

I reduced the

max_length

to 512 and used padding. That seems to have unblocked me and gotten me past that bump in the road. Not sure if that will impact accuracy significantly, but seems to be consistent with some of the tricks being used here: https://docs.ray.io/en/master/ray-air/examples/gptj_deepspeed_fine_tuning.html

Nice, glad you are unblocked!

Thanks so much for your help! Appreciate it!!!

Anytime! Keep us posted if there are other models you are hoping to get running on your Metaflow deployments. Happy to kick the tires and smooth out the path for you as much as we can.

Much obliged. Thank you Eddie!

Ya I think thats right. Btw I see a thread on original Dolly repo on this topic which pointed to a bug in the dolly-metaflow

trainer.py

file with the response key. It is fixed now so think it may work if you can either pull change or do the minimal edit changing

RESPONSE_KEY

to

RESPONSE_KEY_NL

in the

response_token_ids = self.tokenizer.encode(RESPONSE_KEY_NL)

line.

Yup did that. Now removing rows in the training data that have token length > 512 and don't have

### RESPONSE KEY\n

Congrats Riley! Appreciate the persistence flex

Couldn't have done it without your guidance Eddie. Thanks a bunch!

awesome! dancingpanda