import os import json import hashlib import inspect import torch from datasets import load_dataset, Dataset, concatenate_datasets from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig from peft import LoraConfig, get_peft_model from trl import SFTTrainer, SFTConfig torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen3-Coder-30B-A3B-Instruct") OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/workspace/outputs/dinodev-m4-qwen3-coder-30b-code-sft-adapter") MAX_LENGTH = int(os.environ.get("MAX_LENGTH", "2048")) MAX_STEPS = int(os.environ.get("MAX_STEPS", "1200")) MAX_SAMPLES = int(os.environ.get("MAX_SAMPLES", "120000")) print("Base model:", BASE_MODEL) print("Output:", OUTPUT_DIR) print("Max length:", MAX_LENGTH) print("Max steps:", MAX_STEPS) print("Max samples:", MAX_SAMPLES) DATASETS = [ "m-a-p/CodeFeedback-Filtered-Instruction", "theblackcat102/evol-codealpaca-v1", "iamtarun/code_instructions_120k_alpaca", "iamtarun/python_code_instructions_18k_alpaca", ] SYSTEM_PROMPT = ( "You are DinoDev, a senior full-stack coding assistant. " "Write clean, production-ready code. Explain important decisions briefly. " "Prefer secure, maintainable, testable solutions." ) def first_value(row, keys): for k in keys: if k in row and row[k] is not None and str(row[k]).strip(): return str(row[k]).strip() return "" def row_to_messages(row): # Already chat formatted if "messages" in row and isinstance(row["messages"], list): return row["messages"] instruction = first_value(row, [ "instruction", "query", "question", "prompt", "input", "problem", "task" ]) extra_input = first_value(row, [ "context", "additional_input", "given", "description" ]) output = first_value(row, [ "output", "response", "answer", "completion", "solution", "code" ]) # Some alpaca rows have prompt containing full instruction if not instruction and "text" in row: instruction = str(row["text"]).strip() if extra_input and extra_input != instruction: user_content = instruction + "\n\nInput / Context:\n" + extra_input else: user_content = instruction if not user_content or not output: return None return [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_content}, {"role": "assistant", "content": output}, ] def load_and_format(tokenizer): formatted = [] seen = set() for ds_name in DATASETS: print(f"\nLoading dataset: {ds_name}") try: ds = load_dataset(ds_name, split="train") except Exception as e: print(f"Skipping {ds_name}: {e}") continue print(ds) # Shuffle before taking rows ds = ds.shuffle(seed=42) for row in ds: messages = row_to_messages(row) if not messages: continue text = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=False, ) # dedupe h = hashlib.sha256(text.encode("utf-8")).hexdigest() if h in seen: continue seen.add(h) # length safety before tokenization-heavy trainer if len(text) < 100: continue formatted.append({"text": text}) if len(formatted) >= MAX_SAMPLES: break print(f"Collected so far: {len(formatted)}") if len(formatted) >= MAX_SAMPLES: break if not formatted: raise RuntimeError("No training samples created. Check dataset schemas.") return Dataset.from_list(formatted) tokenizer = AutoTokenizer.from_pretrained( BASE_MODEL, trust_remote_code=True, use_fast=True, ) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token dataset = load_and_format(tokenizer) dataset = dataset.train_test_split(test_size=0.02, seed=42) train_ds = dataset["train"] eval_ds = dataset["test"] print("Train rows:", len(train_ds)) print("Eval rows:", len(eval_ds)) print("Sample text:\n", train_ds[0]["text"][:1000]) bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, ) model = AutoModelForCausalLM.from_pretrained( BASE_MODEL, quantization_config=bnb_config, device_map="auto", trust_remote_code=True, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, ) model.config.use_cache = False # Low-VRAM QLoRA prep: # Avoid PEFT prepare_model_for_kbit_training() because it can cast large params to fp32 # and cause CUDA OOM on 30B MoE models. for param in model.parameters(): param.requires_grad = False if hasattr(model, "gradient_checkpointing_enable"): model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) if hasattr(model, "enable_input_require_grads"): model.enable_input_require_grads() else: def make_inputs_require_grad(module, input, output): output.requires_grad_(True) model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) lora_config = LoraConfig( r=16, lora_alpha=32, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM", # Attention-only LoRA first. Safer for 30B MoE / 40GB-80GB GPU. target_modules=[ "q_proj", "k_proj", "v_proj", "o_proj", ], ) model = get_peft_model(model, lora_config) model.print_trainable_parameters() # Make SFTConfig compatible with older/newer TRL cfg_params = inspect.signature(SFTConfig.__init__).parameters kwargs = dict( output_dir=OUTPUT_DIR, per_device_train_batch_size=1, gradient_accumulation_steps=16, learning_rate=2e-4, max_steps=MAX_STEPS, warmup_ratio=0.03, logging_steps=10, save_steps=100, eval_steps=100, bf16=True, optim="paged_adamw_8bit", lr_scheduler_type="cosine", gradient_checkpointing=True, report_to="none", save_total_limit=3, ) if "eval_strategy" in cfg_params: kwargs["eval_strategy"] = "steps" elif "evaluation_strategy" in cfg_params: kwargs["evaluation_strategy"] = "steps" if "dataset_text_field" in cfg_params: kwargs["dataset_text_field"] = "text" if "packing" in cfg_params: kwargs["packing"] = True if "max_seq_length" in cfg_params: kwargs["max_seq_length"] = MAX_LENGTH elif "max_length" in cfg_params: kwargs["max_length"] = MAX_LENGTH training_args = SFTConfig(**kwargs) try: trainer = SFTTrainer( model=model, args=training_args, train_dataset=train_ds, eval_dataset=eval_ds, processing_class=tokenizer, ) except TypeError: trainer = SFTTrainer( model=model, args=training_args, train_dataset=train_ds, eval_dataset=eval_ds, tokenizer=tokenizer, ) trainer.train() trainer.save_model(OUTPUT_DIR) tokenizer.save_pretrained(OUTPUT_DIR) readme = f"""--- base_model: {BASE_MODEL} library_name: peft tags: - qwen3 - coder - qlora - lora - peft - dinodev - code-sft --- # DinoDev M4 Qwen3 Coder 30B Code SFT Adapter Base model: `{BASE_MODEL}` Training type: QLoRA / PEFT LoRA adapter Datasets: - m-a-p/CodeFeedback-Filtered-Instruction - theblackcat102/evol-codealpaca-v1 - iamtarun/code_instructions_120k_alpaca - iamtarun/python_code_instructions_18k_alpaca Settings: - max_length: {MAX_LENGTH} - max_steps: {MAX_STEPS} - max_samples: {MAX_SAMPLES} - LoRA r: 32 - LoRA alpha: 64 """ with open(os.path.join(OUTPUT_DIR, "README.md"), "w") as f: f.write(readme) print("Training complete. Saved to:", OUTPUT_DIR)