[NVIDIA] Update scripts to improve configurability

google / paxml

Pax is a Jax-based machine learning framework for training large scale models. Pax allows for advanced and fully configurable experimentation and parallelization, and has demonstrated industry leading model flop utilization rates.

Apache License 2.0

458 stars 69 forks source link

[NVIDIA] Update scripts to improve configurability #41

Closed ashors1 closed 1 year ago

ashors1 commented 1 year ago

Adds missing 'import math' statement
Switches to Fiddle to allow hparams to be overridden from the command line and updates scripts accordingly