Which dataset can I use to retrain a pruned model

Sixte_Oriol_Llenas_S: <p>Hello,</p>
<p>I am currently pruning the following LLMs:</p>
<ul>
<li>Llama-2-7b-hf</li>
<li>Llama-2-13b-hf</li>
<li>DeepSeek-R1-Distill-Llama-8B</li>
<li>DeepSeek-R1-Distill-Qwen-32B</li>
</ul>
<p>I would like to know which data is the best to retrain these models after the pruning process. I am currently using C4 dataset and I preprocess it as follows:</p>
<pre><code class="lang-auto">    raw_datasets = load_dataset(datasets_path, split=[f'train[0:{train_number_samples}]', f'validation[:{validation_number_samples}]'], cache_dir="./cache_training")
    raw_datasets = DatasetDict({
        "train": raw_datasets[0],
        "validation": raw_datasets[1],
    })

    column_names = list(raw_datasets["train"].features)
    text_column_name = "text" if "text" in column_names else column_names[0]

    def tokenize_function(examples):
        return tokenizer(examples[text_column_name])

    tokenized_datasets = raw_datasets.map(
        tokenize_function,
        batched=True,
        remove_columns=column_names,
    )

    # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size_texts.
    def group_texts(examples):
        block_size_texts = 1024
        # Concatenate all texts.
        concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
        total_length = len(concatenated_examples[list(examples.keys())[0]])
        # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
            # customize this part to your needs.
        total_length = (total_length // block_size_texts) * block_size_texts
        # Split by chunks of max_len.
        result = {
            k: [t[i : i + block_size_texts] for i in range(0, total_length, block_size_texts)]
            for k, t in concatenated_examples.items()
        }
        result["labels"] = result["input_ids"].copy()
        return result

    lm_datasets = tokenized_datasets.map(
        group_texts,
        batched=True,
    )
</code></pre>
<p>I am also using the following training args:</p>
<pre><code class="lang-auto">training_args=transformers.TrainingArguments(
	output_dir="./results",

	#Evaluation
	eval_strategy="steps",
	eval_steps=250,
	eval_accumulation_steps=2,
	eval_on_start=True,

	#Training
	learning_rate=lr,
	num_train_epochs=2,
	per_device_train_batch_size=4,
	gradient_accumulation_steps=4,
	gradient_checkpointing=True,
	max_grad_norm=1.0,

	#Precision
	bf16=True,

	#Logging/Saving
	logging_strategy="steps",
	logging_steps=50,
	report_to="none",
	save_strategy="no",

	#Optimization
	optim="adamw_torch_fused",
	lr_scheduler_type="cosine",
	warmup_ratio=0.1,

	#System
	ddp_find_unused_parameters=False,
	dataloader_pin_memory=True,
	dataloader_num_workers=4
)
</code></pre>
<p>My goal is to get all the metrics back. In particular, I am using lm-eval (<a href="https://github.com/EleutherAI/lm-evaluation-harness" class="inline-onebox" rel="noopener nofollow ugc">GitHub - EleutherAI/lm-evaluation-harness: A framework for few-shot evaluation of language models.</a>) and evaluate the models on <em>arc_challenge_acc, arc_challenge_acc_norm, arc_easy_acc, arc_easy_acc_norm, hellaswag_acc, hellaswag_acc_norm, lambada_openai_acc, lambada_standard_acc, openbookqa_acc, openbookqa_acc_norm, piqa_acc, piqa_acc_norm, winogrande_acc</em>.</p>
<p>Question: is the setup I am using good enough? Would you try any other approach?</p>
<p>Thanks!</p>