axlearn icon indicating copy to clipboard operation
axlearn copied to clipboard

Add Dataflow Inference Examples

Open jiya-zhang opened this issue 1 year ago • 0 comments

This PR adds 2 examples for running batch inference on Dataflow:

  1. Using a Custom Model Handler for JAX models
  2. Using a Built-in HuggingFace Model Handler

These pipelines can run on CPUs or GPUs. In order to run on GPUs, see Link. Users would need to create their own custom image with the necessary libraries, and pass in an additional command line flag.

Future work for this PR:

  1. Implement sharing model across threads on the same machine. See Link.
  2. Write unit tests

jiya-zhang avatar May 15 '24 04:05 jiya-zhang