def compute_gradients(model, data_loader, device):
    model.train()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    loss_function = torch.nn.CrossEntropyLoss()

    batch = next(iter(data_loader))
    inputs = batch["input_ids"].squeeze(1).to(device)
    attention_mask = batch["attention_mask"].squeeze(1).to(device)
    labels = inputs.clone()

    optimizer.zero_grad()
    outputs = model(inputs, attention_mask=attention_mask)
    logits = outputs.logits
    loss = loss_function(logits.view(-1, logits.size(-1)), labels.view(-1))
    loss.backward()

def compute_importance_scores_with_gradients(model):
    importance_scores = {}

    for name, param in model.named_parameters():
        if param.requires_grad and param.grad is not None:
            importance = (param * param.grad).abs()
            importance_scores[name] = importance

    return importance_scores

def prune_mama_mutation_2(args, model, tokenizer, device=torch.device("cuda:0"), prune_n=0, prune_m=0):
    # TODO: Optimize the MAMA pruning algorithm based on advanced indicators.
    # Last Updated Date: 20240921 - 
    # status: under review
    from torch.utils.data import DataLoader
    from datasets import load_dataset

    # Load a small dataset for gradient computation
    dataset = load_dataset('wikitext', 'wikitext-2-raw-v1', split='validation')
    def tokenize_function(examples):
        return tokenizer(examples['text'], return_tensors='pt', truncation=True, padding='max_length', max_length=512)

    tokenized_dataset = dataset.map(tokenize_function, batched=True)
    data_loader = DataLoader(tokenized_dataset, batch_size=1)

    model.to(device)
    compute_gradients(model, data_loader, device)
    importance_scores = compute_importance_scores_with_gradients(model)
    prune_parameters(model, importance_scores, args.sparsity_ratio) # the output log: (prune_llm) [ec2-user@ip-172-31-5-145 suanfamama-prune-1]$ python main.py --model baffo32/decapoda-research-llama-7B-hf --prune_method mama_mutation_2 --sparsity_ratio 0.50 --sparsity_type unstructured --save out/llama_7b/unstructured/prune_mampython main.py --model baffo32/decapoda-research-llama-7B-hf --prune_method mama_mutation_2 --sparsity_ratio 0.50 --sparsity_type unstructured --save out/llama_7b/unstructured/prune_mama_mutation_2/
torch 1.10.1
transformers 4.28.0
accelerate 0.18.0
# of gpus:  1
loading llm model baffo32/decapoda-research-llama-7B-hf
/home/ec2-user/miniconda3/envs/prune_llm/lib/python3.9/site-packages/huggingface_hub/file_download.py:1150: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.
  warnings.warn(
Loading checkpoint shards:   0%|                                                                                                                                                         Loading checkpoint shards:   3%|██████▌                                                                                                                                                  Loading checkpoint shards:   6%|█████████████                                                                                                                                            Loading checkpoint shards:   9%|███████████████████▌                                                                                                                                     Loading checkpoint shards:  12%|██████████████████████████                                                                                                                               Loading checkpoint shards:  15%|████████████████████████████████▌                                                                                                                        Loading checkpoint shards:  18%|███████████████████████████████████████                                                                                                                  Loading checkpoint shards:  21%|█████████████████████████████████████████████▌                                                                                                           Loading checkpoint shards:  24%|████████████████████████████████████████████████████                                                                                                     Loading checkpoint shards:  27%|██████████████████████████████████████████████████████████▋                                                                                              Loading checkpoint shards:  30%|████████████████████████████████████████████████████████████████▊                                                                                        Loading checkpoint shards:  33%|███████████████████████████████████████████████████████████████████████▎                                                                                 Loading checkpoint shards:  36%|█████████████████████████████████████████████████████████████████████████████▊                                                                           Loading checkpoint shards:  39%|████████████████████████████████████████████████████████████████████████████████████▎                                                                    Loading checkpoint shards:  42%|██████████████████████████████████████████████████████████████████████████████████████████▊                                                                                                                           | 14/33 [00:42<00:58,  3.08s/it]

Loading checkpoint shards:  45%|█████████████████████████████████████████████████████████████████████████████████████████████████▎                                                       Loading checkpoint shards:  48%|███████████████████████████████████████████████████████████████████████████████████████████████████████▊                                                 Loading checkpoint shards:  52%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████▏                                          Loading checkpoint shards:  55%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▋                                    Loading checkpoint shards:  58%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏                             Loading checkpoint shards:  61%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▋                       Loading checkpoint shards:  64%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏                Loading checkpoint shards:  67%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▋          Loading checkpoint shards:  70%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏   Loading checkpoint shards:  73%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████Loading checkpoint shards:  76%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████Loading checkpoint shards:  79%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████Loading checkpoint shards:  82%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████Loading checkpoint shards:  85%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████Loading checkpoint shards:  88%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████Loading checkpoint shards:  91%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████Loading checkpoint shards:  94%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████Loading checkpoint shards:  97%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 33/33 [01:41<00:00,  3.09s/it]
/home/ec2-user/miniconda3/envs/prune_llm/lib/python3.9/site-packages/huggingface_hub/file_download.py:1150: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.
  warnings.warn(
use device  cuda:0
pruning starts
Map:   0%|                                                                                                                                                                                                                                            | 0/3760 [00:00<?, ? examples/s]Using pad_token, but it is not set yet.
Map:   0%|                                                                                                                                                                                                                                            | 0/3760 [00:00<?, ? examples/s]
Traceback (most recent call last):
  File "/home/ec2-user/Projects/suanfamama-prune-1/main.py", line 138, in <module>
    main()
  File "/home/ec2-user/Projects/suanfamama-prune-1/main.py", line 83, in main
    prune_mama_mutation_2(args, model, tokenizer, device, prune_n=prune_n, prune_m=prune_m)
  File "/home/ec2-user/Projects/suanfamama-prune-1/lib/prune.py", line 624, in prune_mama_mutation_2
    tokenized_dataset = dataset.map(tokenize_function, batched=True)
  File "/home/ec2-user/miniconda3/envs/prune_llm/lib/python3.9/site-packages/datasets/arrow_dataset.py", line 560, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/home/ec2-user/miniconda3/envs/prune_llm/lib/python3.9/site-packages/datasets/arrow_dataset.py", line 3035, in map
    for rank, done, content in Dataset._map_single(**dataset_kwargs):
  File "/home/ec2-user/miniconda3/envs/prune_llm/lib/python3.9/site-packages/datasets/arrow_dataset.py", line 3438, in _map_single
    batch = apply_function_on_filtered_inputs(
  File "/home/ec2-user/miniconda3/envs/prune_llm/lib/python3.9/site-packages/datasets/arrow_dataset.py", line 3300, in apply_function_on_filtered_inputs
    processed_inputs = function(*fn_args, *additional_args, **fn_kwargs)
  File "/home/ec2-user/Projects/suanfamama-prune-1/lib/prune.py", line 622, in tokenize_function
    return tokenizer(examples['text'], return_tensors='pt', truncation=True, padding='max_length', max_length=512)
  File "/home/ec2-user/miniconda3/envs/prune_llm/lib/python3.9/site-packages/transformers/tokenization_utils_base.py", line 2538, in __call__
    encodings = self._call_one(text=text, text_pair=text_pair, **all_kwargs)
  File "/home/ec2-user/miniconda3/envs/prune_llm/lib/python3.9/site-packages/transformers/tokenization_utils_base.py", line 2624, in _call_one
    return self.batch_encode_plus(
  File "/home/ec2-user/miniconda3/envs/prune_llm/lib/python3.9/site-packages/transformers/tokenization_utils_base.py", line 2806, in batch_encode_plus
    padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies(
  File "/home/ec2-user/miniconda3/envs/prune_llm/lib/python3.9/site-packages/transformers/tokenization_utils_base.py", line 2443, in _get_padding_truncation_strategies
    raise ValueError(
ValueError: Asking to pad but the tokenizer does not have a padding token. Please select a token to use as `pad_token` `(tokenizer.pad_token = tokenizer.eos_token e.g.)` or add a new pad token via `tokenizer.add_special_tokens({'pad_token': '[PAD]'})`.