Ashish Rao comments

Repositories
Issues
Comments

Results 3 comments of


                                            Ashish Rao

Checkpointer memory leak

Hello @sirfz , could you provide a couple more details about the exact configuration and hardware you were running? Specifically: - How much host-side RAM did your machine have? -...

Checkpointer memory leak

Thanks for the info; unfortunately I haven't been able to reproduce the issue. Rerunning with Orbax 0.11.19 on a machine with a 40GB A100 and 125 GiB of host RAM...

Aggregate cross-host data transfers made by jax.device_put

Thanks for reviewing this! I looked into the `pxla.shard_args` code, and I think we will need to handle cross-host transfers separately by directly calling `batched_copy_array_to_devices_with_sharding`. `pxla.shard_args` for jax Arrays eventually...