make checkpointing thread safe #245
Conversation
facebook-github-bot
added
the
CLA Signed
tushar00jain
force-pushed
the
pr245
branch
2 times, most recently
from
5d2f3e4 to
dd0f5fc
Compare
There was a problem hiding this comment.
Is there some other way we can handle this?
Two main concerns:
- cloning has significant memory impact
- send_checkpoint may race with the first inner optimizer step so cloining in _to_cpu may not even be safe
I thought with the DiLoCo implementation we had a second copy of the weights to compute the pseudo gradient. Can we not reuse those for the state_dict transfer?
| out.append(v) | ||
| if isinstance(v, DTensor): | ||
| clone = distribute_tensor( | ||
| v.to_local().clone(), v.device_mesh, v.placements |
There was a problem hiding this comment.
Is v.clone() not sufficient? What does that do? Do we also need special logic for cuda DTensors?
There was a problem hiding this comment.
Is v.clone() not sufficient?
tried it and didn't work. it made an empty dtesnor. took a long time to figure how to clone dtensors
| @@ -278,7 +279,13 @@ def _to_cpu(values: List[T], pin_memory: bool) -> List[T]: | |||
| else: | |||
| out.append(v.cpu()) | |||
| else: | |||
| out.append(v) | |||
| if isinstance(v, DTensor): | |||
There was a problem hiding this comment.
should we rename _to_cpu to _clone_cpu?
the model transfer isn't even controlled by the wrapper in general
thought about blocking everything until checkpoint transfer is complete but that's probably also complicated since the node may never want to transfer checkpoint or use locks, the inner step locks the model and we lock while actually transferring the state dict |
tushar00jain
force-pushed
the
pr245
branch
3 times, most recently
from
0b91490 to
49a161f
Compare
|
@d4l3k for the gpu case, we're clone the tensor into cpu memory anyway. thing is we don't control at what inner step the checkpoint will be sent so the regression test ends up being non deterministic. it changes what model parameters are used for syncing, even if we make it thread safe also on thread safety, the cloning seems to be a blocking call in post step hook, so it shouldn't race with the inner optimizer step? i added some locking anyway |
tushar00jain
force-pushed
the
pr245
branch
3 times, most recently
from
636a86a to
2b7defd
Compare
tushar00jain
changed the title
tushar00jain
force-pushed
the
pr245
branch
6 times, most recently
from
cc0a37a to
76f2f60
Compare
tushar00jain
changed the title
tushar00jain
force-pushed
the
pr245
branch
2 times, most recently
from
2388879 to
fc1fa08
Compare
tushar00jain
force-pushed
the
pr245
branch
2 times, most recently
from
f99b8dd to
acd9ede
Compare
| "user": {key: value() for key, value in self._user_state_dicts.items()}, | ||
| "torchft": self.state_dict(), | ||
| } | ||
| with self._state_dict_lock.r_lock(): |
There was a problem hiding this comment.
we already have allow_checkpoint and disallow_checkpoint in HTTPTransport -- can we reuse those instead?
There was a problem hiding this comment.
- that also requires copying over state dict to the http transport
- or keeping some tracking on which step was transferred last to the http transport
with a separate lock, we can decouple checkpoint specific logic with training logic
tushar00jain
force-pushed
the
pr245
branch
4 times, most recently
from
0283856 to
1f5854d
Compare
tushar00jain
changed the title
tushar00jain
force-pushed
the
pr245
branch
2 times, most recently
from
800c48f to
577bacd
Compare
| @@ -324,6 +328,21 @@ def __init__( | |||
| # first step is 1 | |||
| self._participating_replica_rank: Optional[int] = None | |||
| self._participating_replica_world_size: int = 0 | |||
| self._is_state_dict_read_allowed = True | |||
|
|
|||
| def allow_state_dict_read(self) -> None: | |||
There was a problem hiding this comment.
Can we add specific tests for these methods? Lock logic is pretty risk prone so would be nice to have unit test coverage for these
tushar00jain
force-pushed
the
pr245
branch
3 times, most recently
from
495ab9a to
595e7e9
Compare
Summary: - the checkpointing wasn't thread safe for http transport so use lock the model in the pre step hook and when we want to transfer the checkpoint
Summary: