(argv: Sequence[str])
| 178 | |
| 179 | |
| 180 | def main(argv: Sequence[str]) -> None: |
| 181 | jax.config.update("jax_default_prng_impl", "unsafe_rbg") |
| 182 | os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0" |
| 183 | |
| 184 | config = pyconfig.initialize(argv) |
| 185 | validate_config(config) |
| 186 | max_utils.print_system_information() |
| 187 | |
| 188 | quantization = LayerwiseQuantization(config) |
| 189 | rng = jax.random.PRNGKey(1234) |
| 190 | rng, rng_quant_params = jax.random.split(rng) |
| 191 | |
| 192 | # load_and_quantize will load a checkpoint and quantize if the following parameters are set: |
| 193 | # quantization=$valid_quantization_type \ |
| 194 | # save_quantized_params_path=$gsbucket_path \ |
| 195 | # checkpoint_is_quantized=false (default) |
| 196 | quantization.load_and_quantize(rng_quant_params) |
| 197 | |
| 198 | |
| 199 | def validate_config(config): |
nothing calls this directly
no test coverage detected