Should we have a unittest for this file if possible?

see #155192, this is now added torch.cuda. I think you can remove these changes as that PR has also updated state_dict_utils

Probably not a good idea to ask users to selectively call close(). I would suggest that users should always call close(). Also if this is a public API, the docstring should follow the template. You can check other docstring.

Should we also wait for the last async_save inside this API as well?

I think we should just leave that to the users. In general, I want to limit global training state within DCP as much as possible.

Can we cache the state in AsyncStager and let user manage the lifetime of async stager? this will allow user to init async stager and destroy as needed. this can be done on every checkpoint or at the end of the job as the user sees fit.

Now, close() method is out of context. It is not on any resource/obj. It is hard for users to understand what close means and why they have to call it.

It will be nice that both paths return AsyncSaveResponse -- one path has only upload_future and another path has both. But I understand that this will break BC. Not sure if we can do some tricks, like making AsyncSaveResponse inherit from Future. Just a thought, not necessarily work.

yeah it's unfortunate but don't think there is a clean way to do this :/

We should remove this line. Users should be able to control the logging level, not the module.

see #155192, this is now added torch.cuda. I think you can remove these changes as that PR has also updated state_dict_utils

why change this? if it is not needed, can we revert it?

nit: ret[key] = _iterate_state_dict(...)

unrelated to this PR: but we should deprecate this in favor of async_process_executor.

repeated, remove?

remove? see suggestion on L321

Can we cache the state in AsyncStager and let user manage the lifetime of async stager? this will allow user to init async stager and destroy as needed. this can be done on every checkpoint or at the end of the job as the user sees fit.

Now, close() method is out of context. It is not on any resource/obj. It is hard for users to understand what close means and why they have to call it.

do we need to handle this case? can we ask user to implement async stager if they need to use zero-copy? I think it is simpler to support and cleaner from API point of view.

pin_memory and share_memory needs to be controlled options, so the user has a choice to disable them as they come with drawbacks that might not work for every model or every system.

should this be STATE_DICT_TYPE | Future[STATE_DICT_TYPE]?

nit: i see we are converting the state_dict_type to future to pass in which is okay but i think it is more readable to just keep passing in the Union of state_dict_type and future.

That makes sense. I was originally thinking in the long term, stage would only return a future (would work for sync as well) but I don't see that happening because it would introduce breaking changes and would be hard to cleanly deprecate the old return type

nit:

is save method still used for sync save? why not change it to support union?

nit: what do you think about this? we can eliminate the save_wrapper and add the if instance check in the save method?

is this assert needed? mypy typechecks should catch if you arent returning a future?

Without this, we introduce a linter error because async_save either returns a Tuple of staging_future/upload_future or an upload future now.

clean up/update CheckpointStager? i think these are from dcp evolution work..

i think we do not want users to create a stager each time. stager caches the storages, may be this needs an update.

throw an exception and suggest synchronizing using the future or call staging_future.result() here?