ZeRO 3 NVMe Offload?

[yandex]

ZeRO 3 NVMe Offload?

[pskucherov]

@pskucherov do you have a fork that's ready to pick up and start from for a single GPU?

No, I don't have a working code yet. But it will work soon.
Running on nvme in this config is useless. On CPU RAM it is more realistic, although it is still very long.

In any case, it's all in a private repository.

[MichaelEk]

You can use ZeRO Offload while finetuning your model.
You can extend our version of Megatron by adding CPU/NVME offload for inference, it is not very hard. However, GPU wil have to download entire model (200 GB) for each inferred token. It takes additional 8 sec for each inferred token in best configuration =(

[lostmsu]

@MichaelEk can't you generate multiple tokens in parallel for different string prefixes? E.g. forward with a batch of 4 sentences. That way one can remove the bottleneck of streaming weights to GPU memory if their need permits such use.

[Vbansal21]

@MichaelEk Thanks for replying.

Actually I wasn't able to implement the offload while inference, that's why I created an issue here. It would be nice if an updated script could be provided.😅

[MichaelEk]

@MichaelEk can't you generate multiple tokens in parallel for different string prefixes? E.g. forward with a batch of 4 sentences. That way one can remove the bottleneck of streaming weights to GPU memory if their need permits such use.

Yes, actually you can. If you have a lot input prefixes, you can significantly reduce the the total time spent loading weights. However, the total time of one generation will be too long to use it, for example, in chatbot.

[pskucherov]

You can use ZeRO Offload while finetuning your model. You can extend our version of Megatron by adding CPU/NVME offload for inference, it is not very hard.

Thanks for the hint. Was able to run on RTX 3070 TI 😃

can't you generate multiple tokens in parallel for different string prefixes?
Yes, actually you can. If you have a lot input prefixes, you can significantly reduce the the total time spent loading weights.

Can you tell a little more about this? I feel that now I need to do this, but I do not know where to start.

[lostmsu]

@pskucherov do you have a fork that's ready to pick up and start from for a single GPU?

[alkavan]

@pskucherov Being able to run on RTX 3070 Ti is really great.
If you could share some code how to "ZeRO Offload" that would be nice!

[Vbansal21]

Huggingface accelerate seems to be the solution. Closing issue.