Beaver fix; working towards another model

[natolambert]

Jacob, can you try running this new Ziya model?

[jacob-morrison]

Ziya works if you use --direct_load, and then you can even crank the batch size to 32. I think the standard pipeline abstraction is inefficient with memory, so we should assume we're using --direct_load for >=7B models.

I'll have results on HERM in ~10 minutes, and then I'll run it on pref sets. I can include these changes in our docker image when this is merged.

[natolambert]

@jacob-morrison on what GPU? Quantization should be exactly the same in the pipeline abstraction?
I was debugging on A6000

[jacob-morrison]

@jacob-morrison on what GPU? Quantization should be exactly the same in the pipeline abstraction? I was debugging on A6000

It should be, but it's clearly not based on my testing. I was running this on an A100 and I couldn't get it to successfully run on a single example with: batch_size=1, max_length=1024. After switching to direct_load, I could change back to: batch_size=32, max_length=2048, with no other changes.

We also already use direct_load for all other models of the same size or bigger, and I'd imagine we'd run into the same issue for those if we tried to use the pipeline abstraction instead.

[natolambert]

Good to know, I'll add something to the docs that the quantization doesn't always work right for pipeline thanks a lot!

[jacob-morrison]

LGTM