Function run_clip

quantization/autoclip.py:110–206

(
    model, enc,
    w_bit, q_config,
    n_samples=128, seqlen=1024,
    datasets="pile"
)

@torch.no_grad()
def run_clip(
    model, enc,
    w_bit, q_config,
    n_samples=128, seqlen=1024,
    datasets="pile"
):
    print(f"Using {datasets} dataset to do calibation")
    samples = get_calib_dataset(
              datasets=datasets, tokenizer=enc, n_samples=n_samples, block_size=seqlen)
    samples = torch.cat(samples, dim=0)

    inps = []
    layer_kwargs = {}
    
    layers = get_blocks(model)

    layers[0] = layers[0].cuda()
    move_embed(model, "cuda")

    # get input and kwargs to layer 0
    # with_kwargs is only supported in PyTorch 2.0
    # use this Catcher hack for now
    class Catcher(nn.Module):
        def __init__(self, module):
            super().__init__()
            self.module = module

        def forward(self, inp, **kwargs):
            inps.append(inp)
            layer_kwargs.update(kwargs)
            raise ValueError  # early exit to break later inference

    # patch layer 0 to catch input and kwargs
    layers[0] = Catcher(layers[0])
    try:
        model(samples.to(next(model.parameters()).device))
    except ValueError:  # work with early exit
        pass
    del samples
    layers[0] = layers[0].module  # restore
    inps = inps[0]

    layers[0] = layers[0].cpu()
    move_embed(model, "cpu")
    
    gc.collect()
    torch.cuda.empty_cache()

    clip_results = {
        "clip": [],
    }

    # solve layer by layer
    for i in tqdm(range(len(layers)), desc="Running Asymmetric Clipping..."):
        layer = layers[i]
        layer = layer.cuda()
        named_linears = get_named_linears(layer)