Open
Conversation
Add comprehensive QLoRA quantization support for SFT and RL alignment training. This commit adds: - QLoRA quantization approach document covering formats, layer strategies, and hardware considerations - Type-safe configuration system with YAML and CLI support - Training scripts supporting SFT, GRPO, and DPO methods - End-to-end validation script to verify quantization support - Default configuration with hardware auto-detection - Complete requirements and documentation Addresses Issue #333: quantization formats are guaranteed supported end-to-end. Files: - quantization_support/QLORA_QUANTIZATION_APPROACH.md - quantization_support/qlora_config.py - quantization_support/train_qlora.py - quantization_support/validate_quantization.py - quantization_support/default_config.yaml - quantization_support/README.md - quantization_support/requirements.txt
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
NSR9
requested review from
ShanmugaSuntharam,
kalaiyarasan1988,
pankaj1311,
u6yuvi and
wayoutisin
as code owners
…ting P18: IDFT smoke testing into QLoRA quantization branch
…oRA model-agnostic Deep technical review of the quantization_support codebase identified 15 issues (1 critical, 3 high, 6 medium, 5 low). All have been fixed. A TECHNICAL_ANALYSIS.md document has been added with full architecture diagrams, per-component breakdowns, and the issues catalog. CRITICAL: - idft_loss.py: Add .detach() to gamma so gradients do not flow through the IDFT reweighting factor. Without this, the optimization objective differs from arXiv:2602.12222. SFT/GRPO/DPO paths are unaffected. HIGH: - qlora_config.py: Wire exclude_modules into BitsAndBytesConfig via llm_int8_skip_modules for both 4-bit and 8-bit modes. Replace regex patterns with plain module names covering Llama, Mistral, Qwen, phi, GPT-2, Falcon, and BLOOM architectures. - qlora_config.py + default_config.yaml + idft_smoke_config.yaml: Change LoRA target_modules default from a hardcoded Llama-style list to "all-linear", which uses PEFT's auto-detection of all nn.Linear layers. This makes LoRA fully model-agnostic out of the box. - run_idft_smoke_test.py: Add GPU memory cleanup (del trainer, gc.collect, torch.cuda.empty_cache) in finally blocks between sequential Phase 2 training runs to prevent OOM. MEDIUM: - qlora_config.py from_dict(): Copy nested dicts before .pop() to avoid mutating the caller's dictionary. - qlora_config.py from_args(): Change truthiness checks to `is not None` so falsy CLI values like --lora_r 0 pass through. - evaluate_smoke_test.py: Add _extract_score_from_raw() helper with METRIC_KEYS covering acc, acc_norm, exact_match, pass@1, em, score. Fix print_results_table() to actually display benchmark scores. - idft_trainer.py: Replace internal _idft_step_count with self.state.global_step so diagnostic logging frequency matches real training steps regardless of gradient accumulation. - validate_quantization.py: check_model_loading() now inspects embedding/norm/lm_head parameters for quant_state to verify they were not accidentally quantized. LOW: - phi_diagnostic.py + run_idft_smoke_test.py: Add padding_side="left" to tokenizer for correct causal LM batching. - validate_quantization.py: Replace fragile two-parser config loading with direct QLoRAConfig.from_yaml(). - train_qlora.py: Emit one-time warning when text column is empty in format_sft_dataset(). - requirements.txt: Bump transformers>=4.38.0, peft>=0.8.0, trl>=0.8.0 to match actual API usage. New file: - TECHNICAL_ANALYSIS.md: Comprehensive technical document covering system architecture (with mermaid diagrams), per-component breakdown, IDFT loss mathematical deep-dive, and the full issues catalog with resolution status. Breaking changes: - IDFT training results are not reproducible against prior runs (the old gradient behavior was incorrect per the paper). - LoRA now targets all linear layers by default instead of 7 hardcoded module names; this increases trainable params but works with any model. Users can override via YAML or --lora_target_modules CLI flag.
jha-vikas
requested review from
aiplaybookin,
firekind,
sidrocks and
smitasasindran
as code owners
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Context
Recovered branch rebased onto staging. PR #521 merges IDFT smoke testing into this branch first, then this PR merges everything into staging.
Test plan
🤖 Generated with Claude Code