"sgl-kernel/vscode:/vscode.git/clone" did not exist on "0a56b721d5531d674bd332b4d051a75857033e6c"
  • Sean Naren's avatar
    [feat] Add context manager to FSDP for easier child module wrapping (#446) · f3359550
    Sean Naren authored
    This adds a context manager that assists in making child modules with similar defaults.
    Usage:
    ```
    from fairscale.nn.misc import enable_wrap, wrap
    
    with enable_wrap(**handleful_of_important_params):
        layer_1 = wrap(torch.nn.Linear(5, 5))
        layer_2 = wrap(torch.nn.Linear(5, 5), flatten_parameters=True) # Override parameters if you'd like
    
    # without the context manager, creates Linear layer
    layer_1 = wrap(torch.nn.Linear(5, 5))
    ```
    If not within the FSDP context, this would be a no-op. This makes it easier to annotate layers without having to copy any changes in parameters.
    f3359550
test_wrap.py 6 KB