• 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
__init__.py 0 Bytes