num_blocks = len(removed_state["attn"])

drop_config = ["none"] * num_blocks

for i in range(num_blocks):

if removed_state["attn"][i] and removed_state["mlp"][i]:

drop_config[i] = "attn+mlp"

elif removed_state["attn"][i]:

drop_config[i] = "attn"

elif removed_state["mlp"][i]:

drop_config[i] = "mlp"

if legal_to_drop_path is None:

legal_to_drop = {"attn": [True] * num_blocks, "mlp": [True] * num_blocks}

return legal_to_drop

with open(legal_to_drop_path, "r") as file:

lines = file.readlines()

lines = [line.strip() for line in lines]

assert (

len(lines) == num_blocks

), "Number of blocks in model and legal_to_drop file do not match (If two_consecutive is set, number of blocks should be half of the model)"

legal_to_drop = {"attn": [False] * len(lines), "mlp": [False] * len(lines)}

for i in range(len(lines)):

if lines[i] == "attn+mlp":

legal_to_drop["attn"][i] = True

legal_to_drop["mlp"][i] = True

elif lines[i] == "attn":

legal_to_drop["attn"][i] = True

elif lines[i] == "mlp":

legal_to_drop["mlp"][i] = True

for subblock_type in ["attn", "mlp"]:

for i in range(len(legal_to_drop[subblock_type])):

if not legal_to_drop[subblock_type][i] and removed_state[subblock_type][i]:

return False

removed_state = copy.deepcopy(removed_state)

if drop_two_consecutive: # decompress: duplicate every entry

removed_state["attn"] = [removed_state["attn"][i // 2] for i in range(2 * len(removed_state["attn"]))]

removed_state["mlp"] = [removed_state["mlp"][i // 2] for i in range(2 * len(removed_state["mlp"]))]

for subblock_type in ["attn", "mlp"]:

for j in range(len(removed_state[subblock_type])):

if subblock_type == "attn":

subblock = getattr(layers[j], get_attn_layer_name(model))

else:

subblock = getattr(layers[j], get_mlp_layer_name(model))

if removed_state[subblock_type][j]:

make_dummy_forward(subblock, subblock_type)

else:

sign = 1

if invert_fitness:

sign = -1

if fitness_fn == "ppl":

return sign * compute_perplexity(model, data)

else:

model,

layers,

candidates,

num_survive: int,

calibration_data,

num_tokens: int,

drop_two_consecutive: bool,

invert_fitness: bool,

fitness_fn: str = "ppl",

target_logits: Optional[List[torch.Tensor]] = None,

):

calibration_minibatch = []

minibatch_ids = []

target_logits_minibatch = []

tokens_used = 0

while tokens_used < num_tokens: # generate minibatch with exactly num_tokens tokens

minibatch_id = random.randint(0, len(calibration_data) - 1)

if minibatch_id in minibatch_ids: # avoid duplicates

continue

minibatch_ids.append(minibatch_id)

if tokens_used + calibration_data[minibatch_id].shape[1] > num_tokens:

calibration_minibatch.append(calibration_data[minibatch_id][:, : num_tokens - tokens_used])

if fitness_fn == "kl":

target_logits_minibatch.append(target_logits[minibatch_id][:, : num_tokens - tokens_used])

tokens_used = num_tokens

else:

calibration_minibatch.append(calibration_data[minibatch_id])

if fitness_fn == "kl":

target_logits_minibatch.append(target_logits[minibatch_id])

tokens_used += calibration_data[minibatch_id].shape[1]

if len(target_logits_minibatch) == 0:

target_logits_minibatch = None

fitnesses = []

for candidate in candidates:

load_states(model, layers, candidate, drop_two_consecutive)

fitness = compute_fitness(model, calibration_minibatch, fitness_fn, invert_fitness, target_logits_minibatch)

fitnesses.append(fitness)

# Keep only best

best_ids = np.argsort(fitnesses)[:num_survive]

parser = argparse.ArgumentParser()

# Model params

parser.add_argument(

"--model_name_or_path",

type=str,

required=True,

help="The name or path to the model being pruned",

)

parser.add_argument(

"--tokenizer_name",

type=str,

default=None,

help="The name or path to the tokenizer. By default use model tokenizer.",

)

# Data params

parser.add_argument(

"--calibration_data",

type=str,

required=True,

help="The name or dataset or path used for calibration.",

)

parser.add_argument("--calibration_tokens", type=int, required=True, help="Number of tokens for calibration.")

parser.add_argument(

"--calibration_sequence_length", type=int, required=True, help="Length of calibration sequences."

)

parser.add_argument(

"--eval_datasets",

nargs="+",

type=str,

default=["f 8000 ineweb_edu", "wikitext2", "c4"],

help="Datasets used for evaluation",

)

parser.add_argument("--no_eval", action="store_true", help="Whether to skip evaluation")

parser.add_argument("--eval_every", default=1, type=int, help="Eval every # generations.")

parser.add_argument("--eval_tokens", default=524288, type=int, help="Number of tokens for evaluation.")

parser.add_argument("--eval_sequence_length", default=None, type=int, help="Length of evaluation sequences.")

# Sparsification params

parser.add_argument("--sparsity", type=float, required=True, help="Fraction of layers to drop.")

# Logging params

parser.add_argument("--log_wandb", default=False, action="store_true", help="Whether to log to W&B")

# Evolutionary Search params

parser.add_argument("--fitness_fn", choices=["ppl", "kl"], default="kl", help="Fitness function.")

parser.add_argument("--generations", required=True, type=int, help="Number of generations in evolutionary search")

parser.add_argument("--offspring", type=int, required=True, help="Number of offspring generated in each generation")

parser.add_argument("--population_size", type=int, default=1, help="Population size in evolutionary search")

parser.add_argument(

"--initially_generated",

type=int,

required=True,

help="Number of search points generated in the beginning; fittest are selected for the initial population",

)

parser.add_argument(

"--initial_tokens",

type=int,

required=True,

help="Number of calibration tokens used for the initial generation",

)

parser.add_argument(

"--survivors_per_selection",

type=int,

nargs="+",

required=True,

help="Number of survivors after each stage of selection",

)

parser.add_argument(

"--tokens_per_selection",

type=int,

nargs="+",

required=True,

help="Number of calibration tokens at each stage of selection",

)

# Evolutionary Search ablation params

parser.add_argument(

"--invert_fitness", action="store_true", help="Whether to invert the fitness function (search for worst)"

)

parser.add_argument("--max_mutations", type=int, default=3, help="Maximum number of mutations in offspring")

parser.add_argument(

"--legal_to_drop_path",

type=str,

default=None,

help="Path to legal_to_drop file. A block can only be dropped if it is dropped in legal_to_drop configuration.",

)

parser.add_argument("--drop_entire_block", action="store_true", help="Whether to drop entire block (attn+mlp).")

parser.add_argument(

"--drop_two_consecutive",

action="store_true",

help="Only drop pairs of consecutive blocks (first and second, third and fourth,...). Can only be set when entire blocks are dropped.",

)

# Misc params

parser.add_argument(

"--dtype",

type=str,

default="float16",

choices=["float16", "float32", "bfloat16"],

help="dtype to load the model.",

)

parser.add_argument(

"--attn_implementation",

type=str,

default=None,

choices=["eager", "sdpa", "flash_attention_2"],

help="Attention implementation: eager, sdpa, or flash_attention_2",

)

parser.add_argument("--use_fast_tokenizer", action="store_true", help="Whether to use fast tokenizer.")

parser.add_argument("--seed", default=0, type=int, help="Random seed.")

# Save params

parser.add_argument("--save_dir", type=str, help="Where to save sparse model.")

parser.add_argument("--drop_config_dir", type=str, help="Where to save layer drop config.")

args = parser.parse_args()

args = parse_args()

# Sanity checks

assert len(args.survivors_per_selection) == len(

args.tokens_per_selection

), "Lists for selection survivors and tokens must have same length"

assert args.survivors_per_selection[-1] == args.population_size, "Last stage should have population_size survivor"

if args.drop_two_consecutive:

assert args.drop_entire_block, "Can't drop two consecutive without dropping entire block"

assert args.legal_to_drop_path == None, "Not implemented"

# Get device and dtype

assert torch.cuda.is_available()

print(args.generations)

device = f"cuda"

dtype = getattr(torch, args.dtype)

# Fix seed

fix_seed(args.seed)

# Init W&B logger

if args.log_wandb:

assert has_wandb, "`wandb` not installed, try pip install `wandb`"

wandb.init(config=args)

# Load model

model = AutoModelForCausalLM.from_pretrained(

args.model_name_or_path,

device_map="auto",

low_cpu_mem_usage=True,

torch_dtype=dtype,

attn_implementation=args.attn_implementation,

trust_remote_code=True,

)

print(model.config.model_type)

print(model)

model.config.use_cache = False # do not use cache

# Tokenizer

tokenizer = AutoTokenizer.from_pretrained(

args.tokenizer_name or args.model_name_or_path, use_fast=args.use_fast_tokenizer

)

# Load calibration data

args.calibration_sequence_length = args.calibration_sequence_length or model.config.max_position_embeddings

calibration_data = get_data(

args.calibration_data,

args.calibration_tokens,

args.calibration_sequence_length,

tokenizer,

train=True,

)

# Load evaluation data

args.sequence_length = args.eval_sequence_length or model.config.max_position_embeddings

eval_datasets = []

for eval_dataset_name in args.eval_datasets:

eval_datasets.append(

get_data(

eval_dataset_name,

args.eval_tokens, # ignored for WikiText2 and C4

args.eval_sequence_length,

tokenizer,

train=False,

)

)

target_logits = []

if args.fitness_fn == "kl":

# Compute target logits (calibration)

for i in trange(0, len(calibration_data), desc="Computing target logits (calib)", leave=False):

with torch.no_grad():

target_logits.append(model(calibration_data[i].to(device)).logits.cpu())

layers = get_layers(model)

blocks_to_remove = int(args.sparsity * len(layers))

print(f"Removing {blocks_to_remove} blocks")

total_blocks = len(layers)

if args.drop_two_consecutive:

assert total_blocks % 2 == 0 and blocks_to_remove % 2 == 0, "Number of total and removed blocks must be even"

total_blocks = total_blocks // 2 # view two consecutive blocks as one block

blocks_to_remove = blocks_to_remove // 2

for layer in layers:

dummy_initialize(getattr(layer, get_attn_layer_name(model)))

dummy_initialize(getattr(layer, get_mlp_layer_name(model)))

legal_mask = get_legal_mask(

args.legal_to_drop_path, total_blocks

) # mask of blocks that can be dropped (all blocks by default)

initial_population_candidates = (

[]

) # store initially generated search points (only take fittest for first population)

while len(initial_population_candidates) < args.initially_generated:

removed_state = {"attn": [False] * total_blocks, "mlp": [False] * total_blocks}

attn_legal_ind = [i for i in range(total_blocks) if legal_mask["attn"][i]]

attn_remove_ind = random.sample(attn_legal_ind, blocks_to_remove)

for ind in attn_remove_ind:

removed_state["attn"][ind] = True

mlp_legal_ind = [i for i in range(total_blocks) if legal_mask["mlp"][i]]

mlp_remove_ind = random.sample(mlp_legal_ind, blocks_to_remove)

for ind in mlp_remove_ind:

removed_state["mlp"][ind] = True

if args.drop_entire_block:

removed_state["mlp"] = copy.deepcopy(removed_state["attn"])

if removed_state in initial_population_candidates: # avoid duplicates

continue

if not is_valid_state(removed_state, legal_mask):

continue

initial_population_candidates.append(removed_state)

population, train_fitnesses = selection(

model=model,

layers=layers,

candidates=initial_population_candidates,

num_survive=args.population_size,

calibration_data=calibration_data,

invert_fitness=args.invert_fitness,

drop_two_consecutive=args.drop_two_consecutive,

num_tokens=args.initial_tokens,

fitness_fn=args.fitness_fn,

target_logits=target_logits,

)

log_dict = {}

for gen_id in range(args.generations):

print(f"Generation {gen_id + 1}/{args.generations}")

print(f"Train fitness {train_fitnesses[0]:.2e}")

for parent in population:

print(f"Parent: attn: {[int(ele) for ele in parent['attn']]} mlp: {[int(ele) for ele in parent['mlp']]}")

load_states(model, layers, population[0], args.drop_two_consecutive)

log_dict["train_fitness"] = train_fitnesses[0]

# Evaluate current search point

if gen_id % args.eval_every == 0 and not args.no_eval:

for eval_dataset_name, eval_dataset in zip(args.eval_datasets, eval_datasets):

ppl_eval = compute_perplexity(model, eval_dataset)

print(f"{eval_dataset_name}: {ppl_eval:.2f}")

log_dict[f"ppl_eval/{eval_dataset_name}"] = ppl_eval

full_train_ppl = compute_perplexity(model, calibration_data)

print(f"full train ppl: {full_train_pp 6D40 l:.2e}")

log_dict["full_train_ppl"] = full_train_ppl

if args.log_wandb:

wandb.log(log_dict)

offspring_list = []

# Generate offspring by Mutation

while len(offspring_list) < args.offspring:

offspring = copy.deepcopy(random.choice(population))

# Mutation

num_flips = min(

random.randint(1, args.max_mutations), random.randint(1, args.max_mutations)

) # bias towards lower values

for _ in range(num_flips):

remove_type = random.randint(0, 1) # 0 remove attention, 1 remove mlp

if remove_type == 0:

subblock_type = "attn"

else:

subblock_type = "mlp"

remove_ind = random.randint(0, total_blocks - 1)

while offspring[subblock_type][remove_ind]:

remove_ind = random.randint(0, total_blocks - 1)

add_ind = random.randint(0, total_blocks - 1)

while not offspring[subblock_type][add_ind]:

add_ind = random.randint(0, total_blocks - 1)

offspring[subblock_type][remove_ind] = True

offspring[subblock_type][add_ind] = False

if args.drop_entire_block:

offspring["mlp"] = copy.deepcopy(offspring["attn"])

if offspring in offspring_list or offspring in population: # avoid duplicates

continue

if not is_valid_state(offspring, legal_mask):

continue

offspring_list.append(offspring)

# Selection in multiple steps

for num_survive, num_tokens in zip(args.survivors_per_selection, args.tokens_per_selection):

if num_survive == args.survivors_per_selection[-1]:

for i in range(

len(population)

): # Elitist EA: Add search points in current generation to final selection step

if population[i] not in offspring_list:

offspring_list.append(population[i])

offspring_list, train_fitnesses = selection(

model=model,

layers=layers,

candidates=offspring_list,

num_survive=num_survive,

calibration_data=calibration_data,

drop_two_consecutive=args.drop_two_consecutive,

invert_fitness=args.invert_fitness,

num_tokens=num_tokens,

fitness_fn=args.fitness_fn,

target_logits=target_logits,

)

population = offspring_list

layer_drop_config = get_layer_drop_config(population[0])

if args.drop_config_dir:

os.makedirs(args.drop_config_dir, exist_ok=True)

with open(os.path.join(args.drop_config_dir, "layer_drop_config.txt"), "w") as f:

for line in layer_drop_config:

f.write(line + "\n")

if args.save_dir:

os.makedirs(args.save_dir, exist_ok=True)

# Save model

torch.save(model, os.path.join(args.save_dir, "final_model.pth"))

# Save layer drop config

with open(os.path.join(args.save_dir, "layer_drop_config.txt"), "w") as f:

for line in layer_drop_config:

f.write(line + "\n")

print("Final configuration:")

for line in layer_drop_config:

print(line)

# Final evaluation

for eval_dataset_name, eval_dataset in zip(args.eval_datasets, eval_datasets):

ppl_eval = compute_perplexity(model, eval_dataset)

print(f"{eval_dataset_name}: {ppl_eval:.2f}")

log_dict[f"ppl_eval/{eval_dataset_name}"] = ppl_eval

full_train_ppl = compute_perplexity(model, calibration_data)

print(f"full train ppl: {full_train_ppl:.2e}")

log_dict["full_train_ppl"] = full_train_ppl

if args.log_wandb: