BitNet icon indicating copy to clipboard operation
BitNet copied to clipboard

About 'replace_hf.py'

Open chyoob opened this issue 1 year ago • 2 comments

Hello @kyegomez

In the inference code of huggingface_example.py, it appears that replace_hf is executed, followed immediately by inference. However, upon examining replace_hf.py, I noticed it converts linear layers to bitlinear layers and seems to declare new weights. I'm curious if there's a need for additional code to transfer the original weights to the bitlinear layers.

maybe ... like this?

def replace_linears_in_hf(
    model,
):
    """
    Replaces all instances of nn.Linear in the given model with BitLinear15b.

    Args:
        model (nn.Module): The model to modify.

    Returns:
        None
    """
    for name, module in model.named_children():
        if isinstance(module, nn.Linear):
            # Replace the nn.Linear with BitLinear matching in features and and out_features, and add it to the model
            new_module = BitLinear(in_features=module.in_features, out_features=module.out_features, bias=module.bias is not None)

            with torch.no_grad():
                new_module.weight = module.weight
                if module.bias is not None:
                    new_module.bias = module.bias
            setattr(model, name, new_module)
        else:
            # Recursively apply to child modules
            replace_linears_in_hf(module)

Thanks.

chyoob avatar Feb 29 '24 05:02 chyoob

Stale issue message

github-actions[bot] avatar Apr 29 '24 12:04 github-actions[bot]

@chyoob great idea, submit a pull request pls

kyegomez avatar Apr 29 '24 12:04 kyegomez

Stale issue message

github-actions[bot] avatar Jun 29 '24 12:06 github-actions[bot]