axlearn
axlearn copied to clipboard
Makes BoundedAsyncCheckpointManager control `max_concurrent_gb` by the size of local shards instead of global shards.