simplify streaming diloco #233
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This was referenced Jul 14, 2025
Merged
Merged
facebook-github-bot
added
the
CLA Signed
tushar00jain
requested review from
d4l3k and
H-Huang
and removed request for
d4l3k
tushar00jain
force-pushed
the
pr233
branch
3 times, most recently
from
6069b95 to
3d53cdc
Compare
d4l3k
approved these changes
Jul 15, 2025
| pseudogradient = local_param - self.original_parameters[name].to( | ||
| p.device | ||
| pseudogradient = ( | ||
| self.original_parameters[name].to(p.device) - local_param |
There was a problem hiding this comment.
This is flipped because we don't do 1- below?
There was a problem hiding this comment.
what do you mean by 1-? the outer optimizer will do param = param - pseudo_grad. but loss goes down in the direction of -pseudo_grad. this is pretty much why we were seeing the loss going up earlier i think.
There was a problem hiding this comment.
proof of this:
- let's say
pseudo_grad = new_param - old_paramwherenew_param = old_param - grad - outer optimizer step is
new_param = old_param - pseudo_grad = old_param - (old_param - grad) = old_param + grad - this is incorrect because it should just be
new_param = old_param - grad
| @@ -588,7 +553,11 @@ def __init__( | |||
| if sync_every < len(model_fragments): | |||
| raise ValueError("Only 1 fragment can be syncrhonized at a time") | |||
|
|
|||
| if fragment_sync_delay >= sync_every: | |||
| if sync_every % len(model_fragments) != 0: | |||
| raise ValueError("sync_every must divide the number of fragments") | |||
There was a problem hiding this comment.
This makes sense for now -- we can relax this later if it turns out people want to sync different parts of the model at different rates though that has other significant considerations
There was a problem hiding this comment.
we can still sync at different rates by passing in a different sync_every. will need to make it configurable in torchtitan too.
tushar00jain
force-pushed
the
pr233
branch
2 times, most recently
from
606428d to
5d57f7f
Compare
Summary: - move the training loop to a separate file - convert it into a class so that methods can be overridden without having to duplicate code
Summary: - since we made a simplifying assumption that we will only ever have 1 inflight fragment, we can simplify some of the logic particularly getting rid of the local step in manager state - we'll just use the manager's step to determine which fragment to sync - this also allows us to easily support heterogenous hardware by tuning the sync_every setting that will make slower/faster machines to perform less/more local steps before they sync - we can also perform quorum right before preparing a fragment sync - this easily ensures that all replicas will have the same max step and sync the same fragment - fix some numeric issues - the sign of the pseudogradient - inplace lerp when mixing local and global parameters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary:
Stack created with Sapling. Best reviewed with ReviewStack.