%pip uninstall -y flash-attn && pip install flash-attn==2.7.3
%pip uninstall -y torch && pip install torch==2.5.1

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=False,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

model_dir = "Qwen/Qwen3-8B"
tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
    model_dir,
    quantization_config=bnb_config,   
    device_map="auto",  
    attn_implementation="flash_attention_2",  # Add this line
    torch_dtype=torch.bfloat16,
    trust_remote_code=True             
)

peft_config = LoraConfig(
    lora_alpha=16,                           # Scaling factor for LoRA
    lora_dropout=0.1,                       # Add slight dropout for regularization
    r=64,                                    # Rank of the LoRA update matrices
    bias="none",                             # No bias reparameterization
    task_type="CAUSAL_LM",                   # Task type: Causal Language Modeling
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ]
)

training_arguments = TrainingArguments(
    output_dir="output",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=1,
    optim="paged_adamw_8bit",
    num_train_epochs=2,
    logging_steps=50,
    max_grad_norm=0.5,
    save_steps=500,
    warmup_ratio=0.1,
    lr_scheduler_type="cosine_with_restarts",
    logging_strategy="steps",
    learning_rate=2e-5,
    fp16=False,
    bf16=True,
    group_by_length=True,
    report_to="none",
    seed = 3407,
    gradient_checkpointing=True,
    weight_decay=0.01
)