


import argparse
import gc
import json
import os

import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.utils.deepspeed import DummyOptim, DummyScheduler


A      = 16
A      = 32





def __UpperCAmelCase       (		__A			)						->					Tuple:



       '''simple docstring'''





       return int(x / 2**2_0			)


class     lowercase__  :






  def __enter__(					self      :					Any ):




         """simple docstring"""







         gc.collect()
         torch.cuda.empty_cache()
         torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
         UpperCAmelCase__										=     torch.cuda.memory_allocated()
         return self






  def __exit__(					self      :					Any		,						*_lowercase      :					List[Any] ):




         """simple docstring"""







         gc.collect()
         torch.cuda.empty_cache()
         UpperCAmelCase__										=     torch.cuda.memory_allocated()
         UpperCAmelCase__										=     torch.cuda.max_memory_allocated()
         UpperCAmelCase__										=     bamb(self.end - self.begin )
         UpperCAmelCase__										=     bamb(self.peak - self.begin )
         # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")





def __UpperCAmelCase       (		__A ,						__A = 1_6 ,						__A = "bert-base-cased" ,						__A = 3_2_0 ,						__A = 1_6_0 ,						)						->					Optional[int]:



       '''simple docstring'''





       UpperCAmelCase__										=     AutoTokenizer.from_pretrained(__A			)
       UpperCAmelCase__										=     load_dataset(
           "glue" ,						"mrpc" ,						split={"train": F"""train[:{n_train}]""", "validation": F"""validation[:{n_val}]"""}			)

       def tokenize_function(__A			):
              # max_length=None => use the model max length (it's actually the default)
              UpperCAmelCase__										=     tokenizer(examples["sentence1"] ,						examples["sentence2"] ,						truncation=__A ,						max_length=__A			)
              return outputs

       # Apply the method we just defined to all the examples in all the splits of the dataset
       UpperCAmelCase__										=     datasets.map(
           __A ,						batched=__A ,						remove_columns=["idx", "sentence1", "sentence2"] ,						load_from_cache_file=__A			)

       # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
       # transformers library
       UpperCAmelCase__										=     tokenized_datasets.rename_column("label" ,						"labels"			)

       def collate_fn(__A			):
              # On TPU it's best to pad everything to the same length or training will be very slow.
              if accelerator.distributed_type == DistributedType.TPU:
                     return tokenizer.pad(__A ,						padding="max_length" ,						max_length=1_2_8 ,						return_tensors="pt"			)
              return tokenizer.pad(__A ,						padding="longest" ,						return_tensors="pt"			)

       # Instantiate dataloaders.
       UpperCAmelCase__										=     DataLoader(
           tokenized_datasets["train"] ,						shuffle=__A ,						collate_fn=__A ,						batch_size=__A			)
       UpperCAmelCase__										=     DataLoader(
           tokenized_datasets["validation"] ,						shuffle=__A ,						collate_fn=__A ,						batch_size=__A			)

       return train_dataloader, eval_dataloader





def __UpperCAmelCase       (		__A ,						__A			)						->					Tuple:



       '''simple docstring'''





       UpperCAmelCase__										=     Accelerator()

       # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
       UpperCAmelCase__										=     config["lr"]
       UpperCAmelCase__										=     int(config["num_epochs"]			)
       UpperCAmelCase__										=     int(config["seed"]			)
       UpperCAmelCase__										=     int(config["batch_size"]			)
       UpperCAmelCase__										=     args.model_name_or_path

       set_seed(__A			)
       UpperCAmelCase__						,			UpperCAmelCase__										=     get_dataloaders(__A ,						__A ,						__A ,						args.n_train ,						args.n_val			)

       # Instantiate the model (we build the model here so that the seed also control new weights initialization)
       UpperCAmelCase__										=     AutoModelForSequenceClassification.from_pretrained(__A ,						return_dict=__A			)

       # Instantiate optimizer
       UpperCAmelCase__										=     (
           AdamW
           if accelerator.state.deepspeed_plugin is None
           or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
           else DummyOptim
       )
       UpperCAmelCase__										=     optimizer_cls(params=model.parameters() ,						lr=__A			)

       if accelerator.state.deepspeed_plugin is not None:
              UpperCAmelCase__										=     accelerator.state.deepspeed_plugin.deepspeed_config[
                  "gradient_accumulation_steps"
              ]
       else:
              UpperCAmelCase__										=     1
       UpperCAmelCase__										=     (len(__A			) * num_epochs) // gradient_accumulation_steps

       # Instantiate scheduler
       if (
           accelerator.state.deepspeed_plugin is None
           or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
       ):
              UpperCAmelCase__										=     get_linear_schedule_with_warmup(
                  optimizer=__A ,						num_warmup_steps=0 ,						num_training_steps=__A ,						)
       else:
              UpperCAmelCase__										=     DummyScheduler(__A ,						total_num_steps=__A ,						warmup_num_steps=0			)

       # Prepare everything
       # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
       # prepare method.
       UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__										=     accelerator.prepare(
           __A ,						__A ,						__A ,						__A ,						__A			)

       # We need to keep track of how many total steps we have iterated over
       UpperCAmelCase__										=     0
       # We also need to keep track of the stating epoch so files are named properly
       UpperCAmelCase__										=     0

       # Now we train the model
       UpperCAmelCase__										=     {}
       for epoch in range(__A ,						__A			):
              with TorchTracemalloc() as tracemalloc:
                     model.train()
                     for step, batch in enumerate(__A			):
                            UpperCAmelCase__										=     model(**__A			)
                            UpperCAmelCase__										=     outputs.loss
                            UpperCAmelCase__										=     loss / gradient_accumulation_steps
                            accelerator.backward(__A			)
                            if step % gradient_accumulation_steps == 0:
                                   optimizer.step()
                                   lr_scheduler.step()
                                   optimizer.zero_grad()

                            overall_step += 1

        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
              accelerator.print("Memory before entering the train : {}".format(bamb(tracemalloc.begin			)			)			)
              accelerator.print("Memory consumed at the end of the train (end-begin): {}".format(tracemalloc.used			)			)
              accelerator.print("Peak Memory consumed during the train (max-begin): {}".format(tracemalloc.peaked			)			)
              accelerator.print(
                  "Total Peak Memory consumed during the train (max): {}".format(
                      tracemalloc.peaked + bamb(tracemalloc.begin			)			)			)
              UpperCAmelCase__										=     tracemalloc.peaked + bamb(tracemalloc.begin			)
              if args.peak_memory_upper_bound is not None:
                     assert (
                         train_total_peak_memory[F"""epoch-{epoch}"""] <= args.peak_memory_upper_bound
                     ), "Peak memory usage exceeded the upper bound"

       accelerator.wait_for_everyone()
       if accelerator.is_main_process:
              with open(os.path.join(args.output_dir ,						"peak_memory_utilization.json"			) ,						"w"			) as f:
                     json.dump(__A ,						__A			)





def __UpperCAmelCase       (		)						->					Any:



       '''simple docstring'''





       UpperCAmelCase__										=     argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage."			)
       parser.add_argument(
           "--model_name_or_path" ,						type=__A ,						default="bert-base-cased" ,						help="Path to pretrained model or model identifier from huggingface.co/models." ,						required=__A ,						)
       parser.add_argument(
           "--output_dir" ,						type=__A ,						default="." ,						help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory." ,						)
       parser.add_argument(
           "--peak_memory_upper_bound" ,						type=__A ,						default=__A ,						help="The upper bound of peak memory usage in MB. If set, the training will throw an error if the peak memory usage exceeds this value." ,						)
       parser.add_argument(
           "--n_train" ,						type=__A ,						default=3_2_0 ,						help="Number of training examples to use." ,						)
       parser.add_argument(
           "--n_val" ,						type=__A ,						default=1_6_0 ,						help="Number of validation examples to use." ,						)
       parser.add_argument(
           "--num_epochs" ,						type=__A ,						default=1 ,						help="Number of train epochs." ,						)
       UpperCAmelCase__										=     parser.parse_args()
       UpperCAmelCase__										=     {"lr": 2E-5, "num_epochs": args.num_epochs, "seed": 4_2, "batch_size": 1_6}
       training_function(__A ,						__A			)


if __name__ == "__main__":
    main()









import warnings
from pathlib import Path
from typing import List, Tuple, Union

import fire
from torch import nn

from transformers import AutoModelForSeqaSeqLM, AutoTokenizer, PreTrainedModel
from transformers.utils import logging


A      = logging.get_logger(__name__)





def __UpperCAmelCase       (		__A ,						__A ,						__A			)						->					None:



       '''simple docstring'''





       UpperCAmelCase__										=     nn.ModuleList([src_layers[i] for i in layers_to_copy]			)
       assert len(__A			) == len(__A			), F"""{len(__A			)} != {len(__A			)}"""
       dest_layers.load_state_dict(layers_to_copy.state_dict()			)


A      = {
    # maps  num layers in teacher -> num_layers in student -> which teacher layers to copy.
    # 12: bart, 16: pegasus, 6: marian/Helsinki-NLP
    12: {
        1: [0],  # This says that if the teacher has 12 layers and the student has 1, copy layer 0 of the teacher
        2: [0, 6],
        3: [0, 6, 11],
        4: [0, 4, 8, 11],
        6: [0, 2, 4, 7, 9, 11],
        9: [0, 1, 2, 4, 5, 7, 9, 10, 11],
        12: list(range(12)),
    },
    16: {  # maps  num layers in student -> which teacher layers to copy
        1: [0],
        2: [0, 15],
        3: [0, 8, 15],
        4: [0, 5, 10, 15],
        6: [0, 3, 6, 9, 12, 15],
        8: [0, 2, 4, 6, 8, 10, 12, 15],
        9: [0, 1, 3, 5, 7, 9, 11, 13, 15],
        12: [0, 1, 2, 3, 4, 5, 6, 7, 9, 11, 13, 15],
        16: list(range(16)),
    },
    6: {1: [0], 2: [0, 5], 3: [0, 2, 5], 4: [0, 1, 3, 5], 6: list(range(6))},
}
A      = {
    # maps  num layers in student -> which teacher layers to copy.
    6: {1: [5], 2: [3, 5], 3: [1, 4, 5], 4: [1, 2, 4, 5]},
    12: {1: [11], 2: [5, 11], 3: [3, 7, 11], 6: [1, 3, 5, 8, 10, 11]},
    16: {1: [15], 4: [4, 9, 12, 15], 8: [1, 3, 5, 7, 9, 11, 13, 15]},
}





def __UpperCAmelCase       (		__A ,						__A			)						->					int:



       '''simple docstring'''





       try:
              UpperCAmelCase__										=     LAYERS_TO_COPY[n_teacher][n_student]
              return val
       except KeyError:
              if n_student != n_teacher:
                     warnings.warn(
                         F"""no hardcoded layers to copy for teacher {n_teacher} -> student {n_student}, defaulting to first"""
                         F""" {n_student}"""			)
              return list(range(__A			)			)





def __UpperCAmelCase       (		__A ,						__A			)						->					List[int]:



       '''simple docstring'''





       if n_student > n_teacher:
              raise ValueError(F"""Cannot perform intermediate supervision for student {n_student} > teacher {n_teacher}"""			)
       elif n_teacher == n_student:
              return list(range(__A			)			)
       elif n_student == 1:
              return [n_teacher - 1]
       else:
              return LAYERS_TO_SUPERVISE[n_teacher][n_student]





def __UpperCAmelCase       (		__A ,						__A = "student" ,						__A = None ,						__A = None ,						__A=False ,						__A=None ,						__A=None ,						**__A ,						)						->					Tuple[PreTrainedModel, List[int], List[int]]:



       '''simple docstring'''





       UpperCAmelCase__										=     "encoder_layers and decoder_layers cannot be both None-- you would just have an identical teacher."
       assert (e is not None) or (d is not None), _msg
       if isinstance(__A ,						__A			):
              AutoTokenizer.from_pretrained(__A			).save_pretrained(__A			)  # purely for convenience
              UpperCAmelCase__										=     AutoModelForSeqaSeqLM.from_pretrained(__A			).eval()
       else:
              assert isinstance(__A ,						__A			), F"""teacher must be a model or string got type {type(__A			)}"""
       UpperCAmelCase__										=     teacher.config.to_diff_dict()

       try:
              UpperCAmelCase__						,			UpperCAmelCase__										=     teacher.config.encoder_layers, teacher.config.decoder_layers
              if e is None:
                     UpperCAmelCase__										=     teacher_e
              if d is None:
                     UpperCAmelCase__										=     teacher_d
              init_kwargs.update({"encoder_layers": e, "decoder_layers": d}			)
       except AttributeError:  # T5
              if hasattr(teacher.config ,						"num_encoder_layers"			):
                     UpperCAmelCase__						,			UpperCAmelCase__										=     teacher.config.num_encoder_layers, teacher.config.num_decoder_layers
              else:
                     UpperCAmelCase__						,			UpperCAmelCase__										=     teacher.config.num_layers, teacher.config.num_decoder_layers
              if e is None:
                     UpperCAmelCase__										=     teacher_e
              if d is None:
                     UpperCAmelCase__										=     teacher_d
              if hasattr(teacher.config ,						"num_encoder_layers"			):
                     init_kwargs.update({"num_encoder_layers": e, "num_decoder_layers": d}			)
              else:
                     init_kwargs.update({"num_layers": e, "num_decoder_layers": d}			)

    # Kwargs to instantiate student: teacher kwargs with updated layer numbers + **extra_config_kwargs
       init_kwargs.update(__A			)

       # Copy weights
       UpperCAmelCase__										=     teacher.config_class(**__A			)
       UpperCAmelCase__										=     AutoModelForSeqaSeqLM.from_config(__A			)
       # Start by copying the full teacher state dict this will copy the first N teacher layers to the student.
       UpperCAmelCase__										=     student.load_state_dict(teacher.state_dict() ,						strict=__A			)
       assert info.missing_keys == [], info.missing_keys  # every student key should have a teacher keys.

       if copy_first_teacher_layers:  # Our copying is done. We just log and save
              UpperCAmelCase__						,			UpperCAmelCase__										=     list(range(__A			)			), list(range(__A			)			)
              logger.info(
                  F"""Copied encoder layers {e_layers_to_copy} and decoder layers {d_layers_to_copy}. Saving them to"""
                  F""" {save_path}"""			)
              student.save_pretrained(__A			)
              return student, e_layers_to_copy, d_layers_to_copy

       # Decide which layers of the teacher to copy. Not exactly alternating -- we try to keep first and last layer.
       if e_layers_to_copy is None:
              UpperCAmelCase__										=     pick_layers_to_copy(__A ,						__A			)
       if d_layers_to_copy is None:
              UpperCAmelCase__										=     pick_layers_to_copy(__A ,						__A			)

       try:
              if hasattr(
                  __A ,						"prophetnet"			):  # For ProphetNet, student.model.encoder.layers is called student.prophetnet.encoder.layers
                     copy_layers(teacher.prophetnet.encoder.layers ,						student.prophetnet.encoder.layers ,						__A			)
                     copy_layers(teacher.prophetnet.decoder.layers ,						student.prophetnet.decoder.layers ,						__A			)
              else:
                     copy_layers(teacher.model.encoder.layers ,						student.model.encoder.layers ,						__A			)
                     copy_layers(teacher.model.decoder.layers ,						student.model.decoder.layers ,						__A			)
       except AttributeError:  # For t5, student.model.encoder.layers is called student.encoder.block
              copy_layers(teacher.encoder.block ,						student.encoder.block ,						__A			)
              copy_layers(teacher.decoder.block ,						student.decoder.block ,						__A			)
       logger.info(
           F"""Copied encoder layers {e_layers_to_copy} and decoder layers {d_layers_to_copy}. Saving them to {save_path}"""			)
       UpperCAmelCase__										=     {
           "teacher_type": teacher.config.model_type,
           "copied_encoder_layers": e_layers_to_copy,
           "copied_decoder_layers": d_layers_to_copy,
       }
       student.save_pretrained(__A			)
       # Save information about copying for easier reproducibility

       return student, e_layers_to_copy, d_layers_to_copy


if __name__ == "__main__":
    fire.Fire(create_student_by_copying_alternating_layers)





