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

Added new config to nvidia.py #79

Closed abhinavgoel95 closed 7 months ago

abhinavgoel95 commented 7 months ago

Related to https://github.com/google/praxis/pull/58

cc @zhangqiaorjc