Maggie Zhang

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...

Dataflow changes

2

I changed a few things in this PR: 1. Add Dockerfile entrypoint for Dataflow (this is needed for Dataflow worker to start up successfully) 2. Mount gcloud config folder to...

Adding demo code as a skeleton to integrate GoodPut measurement into AXLearn. To do's: 1. Consider localizing GoodPut measurement to within trainer.py 2. Make the measurement configurable 3. Automatically pick...

At this time, you need to use Orbax Checkpointing to test out in-memory checkpointing. It is also better integrated with Pathways, when the time to test that comes. This PR...