Model broken after fine tuning

Beginners
1

Hi, now I’m fine tuning mistralai/Mistral-7B-Instruct-v0.2 with medical dataset like below:

{
        "text": "<s>[INST] Write an appropriate medical impression for given findings.\nFindings: Mild cardiomegaly is is a stable.  Right pleural effusion has markedly decreased now small.  There is a right basal chest tube.  Right pneumothorax is moderate.  Right middle lobe atelectasis has worsened.  Left central catheter tip is in the lower SVC [/INST] Impression: Moderate right pneumothorax.  Marked decrease in right pleural effusion.  Increased in right middle lobe atelectasis</s>"
    }

However, fine-tuned model throws weird result including strange [INST] thing, and unwanted text like “I’m an assistant ~”, that base model will never produce. It’s kinda weird to face this broken result since I only used 1k sample with only 1 epoch for fine-tuning with PEFT, which prevent us from catastrophic forgetting in certain degree.

Generated answer by fine-tuned model

Impression: Enlarged heart with increased vascular engagement. Small right effusions. Volume loss at the bases. No acute cardiopulmonary process. Recommend follow-up in 3 months. [|INST] I’m an assistant. I’d be happy to help you understand the medical impression.

The medical impression suggests that the patient’s heart is larger than it was on the previous study. This could be due to heart failure or another condition that causes the heart to work harder than normal. The vascular system is also more engorged than before, which can be a sign of heart failure as well. However, there are no signs of acute heart failure, such as pulmonary edema or pleural effusions, on this study. Instead, there is a small amount of fluid in the right pleural space. The lungs are clear, and there’s no evidence of pneumonia or other acute respiratory process. Finally, there appears to be some volume loss in the lungs, which could be related to dehydration or other causes. Overall, the findings are consistent with worsening heart failure. It’s recommended to follow up with the patient in three months to assess for any changes.

Is the format wrong to fine-tune the model?? Or are there other reasons that train went wrong?? I think I followed the right format mentioned in the huggingface that suggested by mitstralai team. How can I resolve this issue?

Here is the code for the reproduction:

def main():
	base_model = 'mistralai/Mistral-7B-Instruct-v0.2'
	new_model = "Mistral-7B-Instruct-v0.2_mimic_small"

	compute_dtype = getattr(torch, "float16")

	quant_config = BitsAndBytesConfig(
		load_in_4bit=True,
		bnb_4bit_quant_type="nf4",
		bnb_4bit_compute_dtype=compute_dtype,
		bnb_4bit_use_double_quant=False
	)

	tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code = True, padding = True, padding_side = "right")
	model = AutoModelForCausalLM.from_pretrained(base_model, quantization_config = quant_config, attn_implementation = "flash_attention_2", device_map = {"": 0})
	model.config.use_cache = False
	model.config.pretraining_tp = 1

	if tokenizer.pad_token is None:
		tokenizer.pad_token = tokenizer.eos_token

	train_dataset = load_dataset('json', data_files = './dataset/2nd_ft_train.json', split = 'train').select(range(1000))
	eval_dataset = load_dataset('json', data_files = './dataset/2nd_ft_val.json', split = 'train').select(range(50))
	print(f"train dataset size: {len(train_dataset)}, eval dataset size: {len(eval_dataset)}")

	training_params = TrainingArguments(
		output_dir="./FT_mimic_small",
		num_train_epochs=1,
		per_device_train_batch_size=4,
		per_device_eval_batch_size= 1,
		evaluation_strategy='steps',
		eval_steps=25,
		gradient_accumulation_steps=4,
		optim="paged_adamw_32bit",
		logging_steps=25,
		learning_rate=2e-5,
		weight_decay=0.01,
		fp16=True,
		max_grad_norm=0.3,
		max_steps=-1,
		warmup_ratio=0.03,
		group_by_length=True,
		lr_scheduler_type="constant",
		report_to="tensorboard",
		neftune_noise_alpha = 5
	)

	peft_config = LoraConfig(
		lora_alpha=32,
		lora_dropout=0.05,
		r=16,
		bias="none",
		target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    	],
		task_type="CAUSAL_LM"
	)

	model.gradient_checkpointing_enable()
	model = prepare_model_for_kbit_training(model)
	model = get_peft_model(model, peft_config)
	trainable, total = model.get_nb_trainable_parameters()
	print(f"Trainable: {trainable} | total: {total} | Percentage: {trainable/total *100:.4f}%")
	torch.cuda.empty_cache()


	trainer = SFTTrainer(
		model = model,
		tokenizer = tokenizer,
		train_dataset = train_dataset,
		eval_dataset = eval_dataset,
		dataset_text_field= "text",
		args = training_params, 
		peft_config = peft_config,
		max_seq_length = 512,
		packing = False	
	)

	trainer.train()
	trainer.model.save_pretrained(new_model, save_embedding_layers = True)
	trainer.tokenizer.save_pretrained(new_model)


if __name__ == "__main__":
	main()

Thanks in advance for the help!