GKD

[qgallouedec]

No description provided.

[qgallouedec]

from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import GKDTrainer, GKDConfig

tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5")

batch_size = 4
gradient_accumulation_steps = 2
output_dir = f"GKD-bsz{batch_size}-grad_acc{gradient_accumulation_steps}-fixed"


model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
model = AutoModelForCausalLM.from_pretrained(model_id)
teacher_model = AutoModelForCausalLM.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token


training_args = GKDConfig(
    output_dir=output_dir,
    per_device_train_batch_size=batch_size,
    dataloader_drop_last=True,
    gradient_accumulation_steps=gradient_accumulation_steps,
    logging_steps=2, 
)
dummy_dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling")

trainer = GKDTrainer(
    model=model_id,
    teacher_model=model_id,
    args=training_args,
    train_dataset=dummy_dataset["train"],
    processing_class=tokenizer,
)

trainer.train()

I'm not sure how to show that the problem arises with GKD. What I do know is that it should happen. @kashif can you take a look?