Open
Conversation
github-actions
bot
added
model
clessig
reviewed
Mar 15, 2026
| # reshard_after_forward=False keeps FE parameters unsharded | ||
| # during the multi-step rollout loop. | ||
| # Needed for pushforward trick. | ||
| fully_shard(module, reshard_after_forward=False, **fsdp_kwargs) |
Collaborator
There was a problem hiding this comment.
@sophie-xhonneux : is this maybe related to the problem we are seeing with the EMATeacher where we need to reshard?
clessig
reviewed
Mar 15, 2026
SavvasMel
force-pushed
the
SavvasMel/develop/pushf_trick
branch
from
056d4be to
8ad0cff
Compare
github-actions
bot
added
data
SavvasMel
force-pushed
the
SavvasMel/develop/pushf_trick
branch
from
2f89054 to
02a0d46
Compare
clessig
reviewed
Mar 29, 2026
Collaborator
clessig
left a comment
clessig
left a comment
There was a problem hiding this comment.
Essentially ready to merge, just some minor details.
|
|
||
| def __init__(self): | ||
| super().__init__() | ||
| self.fe_blocks = torch.nn.ModuleList() |
Collaborator
There was a problem hiding this comment.
self.blocks = ... Also do we need this as all? If anything I would expect self.blocks = torch.nn.Identity() but since we implement forward it might not be needed
Contributor
Author
There was a problem hiding this comment.
I actually added that for this:
num_params_fe = get_num_parameters(self.forecast_engine.fe_blocks)
See below.
| tokens = tokens.reshape(shape).sum(axis=1) | ||
|
|
||
| # Allow for pushforward trick | ||
| p_fwd = self.cf.get("pushforward_trick", False) |
Collaborator
There was a problem hiding this comment.
I would expect that the push-forward trick config is part of forecast, i.e.
training_config :
....
forecast:
num_steps: 3
push_forward: True
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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.
Description
This PR adds the necessary changes for the pushforward trick.
Issue Number
Closes #1740
Is this PR a draft? Mark it as draft.
Checklist before asking for review
./scripts/actions.sh lint./scripts/actions.sh unit-test./scripts/actions.sh integration-testlaunch-slurm.py --time 60