Add option to FSDP wrap by groups of blocks #340
Conversation
| if recurse: | ||
| return True # always recurse for simplicity |
There was a problem hiding this comment.
I don't understand why this works. When does this ever get called with recurse == False?
There was a problem hiding this comment.
From what I understand this function basically gets called twice on every module. Once with recurse=False to check if that module itself should be wrapped, and once again with recurse=True to check if it should go deeper into submodules of the current module to potentially wrap any of those.
There was a problem hiding this comment.
So the meaning of the return value changes depending on the value of recurse, which is confusing... and is the reason we had the wrapping bug in the first place where we thought we were wrapping by block but we were actually just wrapping the whole model.
| return isinstance(module, OlmoBlock) | ||
|
|
||
| return fsdp_wrap_fn | ||
| elif wrap_strategy == FSDPWrapStrategy.by_block_group: |
There was a problem hiding this comment.
With this strategy, the input and output embeddings are never wrapped. I think that's fine at this point in time, but we should experiment with it.
There was a problem hiding this comment.
Agreed. I didn't want to change too many things.
olmo/config.py
Outdated
| block_group_size: int = 1 | ||
| """ | ||
| The number of blocks to group together into a single parent block. | ||
| This is has no affect on the number of parameters in the model and is only used to wrap groups |
Closes #289.
Adds the ability to wrap multiple sequential transformer blocks together into a single parent FSDP wrapper by setting
--model.block_group_sizeto the desired size and--fsdp.wrapping_strategyto "by_block_group". For example:In this case every group of 4 sequential blocks would be wrapped together.