Function prune_wanda_v2

lib/model_wrapper.py:143–188

(
    args,
    model,
    tokenizer,
    model_base=None,
    device=torch.device("cuda:0"),
    prune_n=0,
    prune_m=0,
    prune_data="wikitext",
)

def prune_wanda_v2(
    args,
    model,
    tokenizer,
    model_base=None,
    device=torch.device("cuda:0"),
    prune_n=0,
    prune_m=0,
    prune_data="wikitext",
):
    model = make_Act(model, verbose=False)

    print(f"loading calibdation data {prune_data}")
    assert prune_data in [
        "wikitext",
        "alpaca",
        "alpaca_cleaned",
        "alpaca_cleaned_no_safety",
        "align",
        "align_short",
        "misalign",
    ]
    dataloader, _ = get_loaders(
        prune_data,
        nsamples=args.nsamples,
        seed=args.seed,
        seqlen=model.seqlen,
        tokenizer=tokenizer,
        disentangle=args.disentangle,
    )
    print("dataset loading complete")

    clear_act_buffer(model)
    with torch.no_grad():
        for batch in dataloader:
            inp, tar = batch[0].to(device), batch[1].to(device)

            if args.disentangle:
                mask = tar.ne(-100)
                with set_mask(model, mask):
                    model(inp)
            else:
                model(inp)

    _prune_core(args, model, model_base, prune_n, prune_m, prune_mode="activation")
    model = revert_Act_to_Linear(model)


def prune_wandg_v1(