Function prune_wandg_v1

lib/model_wrapper.py:191–238

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

def prune_wandg_v1(
    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")

    model.requires_grad_(False)
    for name, module in model.named_modules():
        if isinstance(module, ActLinear):
            module.base.requires_grad_(True)
            module.base.zero_grad()

    for batch in dataloader:
        inp, tar = batch[0].to(device), batch[1].to(device)
        assert args.disentangle, "should run in disentangle mode"
        with no_act_recording(model):
            loss = model(input_ids=inp, labels=tar)[0]
        loss.backward()

    _prune_core(args, model, model_base, prune_n, prune_m, prune_mode="gradient")
    model = revert_Act_to_Linear(model)
    model.zero_grad()  # freeze gradient to save cuda memory


def prune_wandg(