axlearn
axlearn copied to clipboard
An Extensible Deep Learning Library
# Changes ## Typo fixes Tool-assisted (via ` typos --format brief --write-changes **/*.py` via [typos-cli](https://github.com/crate-ci/typos)). The rest of the effort is fine-tooth combing that output. Aside: If you attempt this...
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...
For Debugging purpose.
Adding a new model configuration for text experiments. The goal is to get an early termination model for fuji-test to accelerate infrastructure validation. + @jiya-zhang
Gradient accumulation allows training with higher batch sizes without scaling out. Added a new learner type ```learner.klass: 'axlearn.common.learner.AccumulatedLearner'``` At a high level the optimization does the following: 1. Input batch...
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...
This is helpful in cases where there is variable step time and looking at the logs would quickly allow you to identify such cases.
Enable unroll for GPU for better communication interleaving