search

elixir-nx / nx

Multi-dimensional arrays (tensors) and numerical definitions for Elixir
2.66k stars 193 forks source link

Allow on_host/2 to execute certain operations on host #1031

Open josevalim opened 1 year ago

josevalim commented 1 year ago

For CPU offloading during training: https://huggingface.co/docs/transformers/main_classes/deepspeed

josevalim commented 1 month ago

Btw, we would implement this with infeed, but infeed is not supported in neither iree or cuda pjrt plugins at the moment. There may be a chance this needs to be implemented in a layer above (such as Axon).