| 1610 | |
| 1611 | |
| 1612 | def prune_wandg_set_difference( |
| 1613 | args, |
| 1614 | model, |
| 1615 | tokenizer, |
| 1616 | model_base=None, |
| 1617 | device=torch.device("cuda:0"), |
| 1618 | prune_n=0, |
| 1619 | prune_m=0, |
| 1620 | prune_data="align_short", |
| 1621 | p=0.5, |
| 1622 | q=0.5, |
| 1623 | ): |
| 1624 | use_cache = model.config.use_cache |
| 1625 | model.config.use_cache = False |
| 1626 | layers = model.model.layers |
| 1627 | if args.use_diff or args.recover_from_base: |
| 1628 | assert model_base is not None |
| 1629 | layers_base = model_base.model.layers |
| 1630 | metric1 = "alpaca_cleaned_no_safety" |
| 1631 | metric2 = prune_data |
| 1632 | |
| 1633 | print( |
| 1634 | "prune p = {}, q = {}, with metric1 = {}, metric2 = {}".format( |
| 1635 | p, q, metric1, metric2 |
| 1636 | ) |
| 1637 | ) |
| 1638 | if args.prune_part: |
| 1639 | print("only prune the layer with low jaccard index") |
| 1640 | else: |
| 1641 | print("prune every linear layer") |
| 1642 | for i in range(len(layers)): |
| 1643 | layer = layers[i] |
| 1644 | subset = find_layers(layer) |
| 1645 | if args.use_diff or args.recover_from_base: |
| 1646 | subset_base = find_layers(layers_base[i]) |
| 1647 | |
| 1648 | if not args.prune_part: |
| 1649 | for name in subset: |
| 1650 | print(f"pruning layer {i} name {name}") |
| 1651 | if args.model == "llama2-7b-chat-hf": |
| 1652 | W_metric1 = pickle.load( |
| 1653 | open( |
| 1654 | f"out/llama2-7b-chat-hf/unstructured/wandg/{metric1}/wanda_score/W_metric_layer_{i}_name_model.layers.{i}.{name}_weight.pkl", |
| 1655 | "rb", |
| 1656 | ) |
| 1657 | ) |
| 1658 | W_metric2 = pickle.load( |
| 1659 | open( |
| 1660 | f"out/llama2-7b-chat-hf/unstructured/wandg/{metric2}/wanda_score/W_metric_layer_{i}_name_model.layers.{i}.{name}_weight.pkl", |
| 1661 | "rb", |
| 1662 | ) |
| 1663 | ) |
| 1664 | elif args.model == "llama2-13b-chat-hf": |
| 1665 | W_metric1 = pickle.load( |
| 1666 | open( |
| 1667 | f"out/llama2-13b-chat-hf/unstructured/wandg/{metric1}/wanda_score/W_metric_layer_{i}_name_model.layers.{i}.{name}_weight.pkl", |
| 1668 | "rb", |
| 1669 | ) |