Skip to content

Commit 6ab852a

Browse files
subramensubramen
committed
add code for measuring skills transferability across LoRA and FFT finetunes
1 parent 0e00335 commit 6ab852a

25 files changed

Lines changed: 112015 additions & 0 deletions

‎end-to-end-use-cases/transferability/README.md‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
# Transferability Research Tool
2+
3+
A Python package for evaluating model transferability across vision-language tasks through systematic fine-tuning and evaluation.
4+
5+
## Directory Structure
6+
7+
```
8+
./
9+
├── config.yaml # Main configuration file
10+
├── experiments/ # Output directory for all experiments
11+
│ └── <experiment_name>/
12+
│ ├── formatted_datasets/ # Processed datasets ready for training
13+
│ ├── finetuned_checkpoints/ # Fine-tuned model checkpoints
14+
│ ├── finetune_logs/ # Training logs
15+
│ ├── grader_logs/ # Evaluation logs per model
16+
│ └── eval_grid_results.json # Final evaluation results
17+
└── transferability/ # Source code package
18+
├── __init__.py # Package entry points
19+
├── __main__.py # Main CLI entry point
20+
├── data/ # Dataset processing
21+
│ ├── __init__.py
22+
│ ├── __main__.py # Module CLI entry point
23+
│ └── dataset_builder.py
24+
├── datasets/ # Dataset format utilities
25+
│ ├── __init__.py
26+
│ └── torchtune_format.py # TorchTune dataset format
27+
├── evals/ # Evaluation utilities
28+
│ ├── __init__.py
29+
│ ├── __main__.py # Module CLI entry point
30+
│ ├── eval_grid.py # Main evaluation grid runner
31+
│ ├── grader.py # Task-specific graders
32+
│ ├── inference.py # Model inference utilities
33+
│ ├── json_grading_utils.py # JSON grading utilities
34+
│ └── shift_analysis.py # Distribution shift analysis
35+
├── finetune/ # Fine-tuning utilities
36+
│ ├── __init__.py
37+
│ ├── __main__.py # Module CLI entry point
38+
│ ├── finetune_grid.py # Main fine-tuning grid runner
39+
│ ├── 8b_full.yaml # TorchTune config for full fine-tuning
40+
│ └── 8b_lora.yaml # TorchTune config for LoRA fine-tuning
41+
└── utils.py # Shared utilities
42+
```
43+
44+
## Usage
45+
46+
Run individual components as Python modules:
47+
48+
```bash
49+
# Prepare datasets
50+
python -m transferability.data ./experiments/my_experiment
51+
52+
# Run fine-tuning grid
53+
python -m transferability.finetune ./experiments/my_experiment
54+
55+
# Run evaluation grid
56+
python -m transferability.evals ./experiments/my_experiment
57+
```
58+
59+
60+
## Configuration
61+
62+
Edit `config.yaml` to configure your tasks, datasets, and training parameters:
63+
64+
```yaml
65+
task1:
66+
dataset: your/huggingface/dataset
67+
system_prompt: "Your system prompt"
68+
user_prompt: "Your user prompt"
69+
image_column: image
70+
assistant_text_column: ground_truth
71+
grader: JSONGrader
72+
sample_percent: 0.01
73+
74+
task2:
75+
# Similar structure for second task
76+
77+
finetuning:
78+
model_path: /path/to/your/base/model
79+
tokenizer_path: /path/to/tokenizer
80+
epochs: 1
81+
batch_size: 8
82+
# Fine-tuning strategy flags
83+
fusion: false
84+
fusion+encoder: false
85+
fusion+decoder: false
86+
fusion+encoder+decoder: true
87+
lora_ranks: [8, 16, 32]
88+
89+
evals:
90+
nb_eval_samples: null # null = use all samples
91+
checkpoint_to_eval: -1 # -1 = use latest checkpoint
92+
model_server_args:
93+
tensor_parallel_size: 2
94+
max_model_len: 4096
95+
```
96+
97+
## Workflow
98+
99+
1. **Configure**: Edit `config.yaml` with your tasks and model paths
100+
2. **Prepare Data**: Download and format datasets from HuggingFace
101+
3. **Fine-tune**: Train models using different strategies (LoRA, full fine-tuning)
102+
4. **Evaluate**: Test all models on all tasks and generate results
103+
104+
## Key Features
105+
106+
- **Modular Design**: Each component can be run independently
107+
- **Multiple Execution Methods**: Module-level, package-level, or direct imports
108+
- **Configurable Tasks**: Define tasks via YAML configuration
109+
- **Grid Search**: Automatically train multiple model variants
110+
- **Comprehensive Evaluation**: Test transferability across tasks
111+
- **Rich Logging**: Detailed logs and metrics for analysis
112+
113+
## Output Structure
114+
115+
Each experiment creates:
116+
- `formatted_datasets/`: HuggingFace datasets converted to training format
117+
- `finetuned_checkpoints/`: Model checkpoints for each training configuration
118+
- `finetune_logs/`: Training metrics and logs
119+
- `grader_logs/`: Per-model evaluation details
120+
- `eval_grid_results.json`: Summary of all evaluation results
121+
122+
## Next Steps
123+
124+
The package is now properly structured for module execution. You can:
125+
126+
1. Update hardcoded paths in `__main__` sections (as planned)
127+
2. Add more sophisticated CLI argument parsing
128+
3. Add configuration validation
129+
4. Add progress tracking and resumption capabilities
130+
5. Add visualization utilities for results analysis

‎end-to-end-use-cases/transferability/config.yaml‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
task1:
2+
dataset: singhsays/fake-w2-us-tax-form-dataset
3+
is_local: false
4+
system_prompt: null
5+
user_prompt: "You are an expert document information extraction system. I will show you an image of a W-2 tax form. Please extract all the information from this form and return it in a JSON format. Include all fields such as employee details, employer details, wages, federal income tax withheld, social security wages, social security tax withheld, medicare wages and tips, medicare tax withheld, and any other information present on the form. Return ONLY the JSON output without any additional text or explanations following this schema {'properties': {'box_b_employer_identification_number': {'title': 'Box B Employer Identification Number', 'type': 'string'}, 'box_c_employer_name': {'title': 'Box C Employer Name', 'type': 'string'}, 'box_c_employer_street_address': {'title': 'Box C Employer Street Address', 'type': 'string'}, 'box_c_employer_city_state_zip': {'title': 'Box C Employer City State Zip', 'type': 'string'}, 'box_a_employee_ssn': {'title': 'Box A Employee Ssn', 'type': 'string'}, 'box_e_employee_name': {'title': 'Box E Employee Name', 'type': 'string'}, 'box_e_employee_street_address': {'title': 'Box E Employee Street Address', 'type': 'string'}, 'box_e_employee_city_state_zip': {'title': 'Box E Employee City State Zip', 'type': 'string'}, 'box_d_control_number': {'title': 'Box D Control Number', 'type': 'integer'}, 'box_1_wages': {'title': 'Box 1 Wages', 'type': 'number'}, 'box_2_federal_tax_withheld': {'title': 'Box 2 Federal Tax Withheld', 'type': 'number'}, 'box_3_social_security_wages': {'title': 'Box 3 Social Security Wages', 'type': 'number'}, 'box_4_social_security_tax_withheld': {'title': 'Box 4 Social Security Tax Withheld', 'type': 'number'}, 'box_5_medicare_wages': {'title': 'Box 5 Medicare Wages', 'type': 'number'}, 'box_6_medicare_wages_tax_withheld': {'title': 'Box 6 Medicare Wages Tax Withheld', 'type': 'number'}, 'box_7_social_security_tips': {'title': 'Box 7 Social Security Tips', 'type': 'number'}, 'box_8_allocated_tips': {'title': 'Box 8 Allocated Tips', 'type': 'number'}, 'box_9_advance_eic_payment': {'anyOf': [{'type': 'string'}, {'type': 'null'}], 'title': 'Box 9 Advance Eic Payment'}, 'box_10_dependent_care_benefits': {'title': 'Box 10 Dependent Care Benefits', 'type': 'number'}, 'box_11_nonqualified_plans': {'title': 'Box 11 Nonqualified Plans', 'type': 'number'}, 'box_12a_code': {'title': 'Box 12A Code', 'type': 'string'}, 'box_12a_value': {'title': 'Box 12A Value', 'type': 'number'}, 'box_12b_code': {'title': 'Box 12B Code', 'type': 'string'}, 'box_12b_value': {'title': 'Box 12B Value', 'type': 'number'}, 'box_12c_code': {'title': 'Box 12C Code', 'type': 'string'}, 'box_12c_value': {'title': 'Box 12C Value', 'type': 'number'}, 'box_12d_code': {'anyOf': [{'type': 'string'}, {'type': 'null'}], 'title': 'Box 12D Code'}, 'box_12d_value': {'title': 'Box 12D Value', 'type': 'number'}, 'box_13_statutary_employee': {'anyOf': [{'type': 'string'}, {'type': 'null'}], 'title': 'Box 13 Statutary Employee'}, 'box_13_retirement_plan': {'anyOf': [{'type': 'string'}, {'type': 'null'}], 'title': 'Box 13 Retirement Plan'}, 'box_13_third_part_sick_pay': {'anyOf': [{'type': 'string'}, {'type': 'null'}], 'title': 'Box 13 Third Part Sick Pay'}, 'box_15_1_state': {'title': 'Box 15 1 State', 'type': 'string'}, 'box_15_1_employee_state_id': {'title': 'Box 15 1 Employee State Id', 'type': 'string'}, 'box_16_1_state_wages': {'title': 'Box 16 1 State Wages', 'type': 'number'}, 'box_17_1_state_income_tax': {'title': 'Box 17 1 State Income Tax', 'type': 'number'}, 'box_18_1_local_wages': {'title': 'Box 18 1 Local Wages', 'type': 'number'}, 'box_19_1_local_income_tax': {'title': 'Box 19 1 Local Income Tax', 'type': 'number'}, 'box_20_1_locality': {'title': 'Box 20 1 Locality', 'type': 'string'}, 'box_15_2_state': {'title': 'Box 15 2 State', 'type': 'string'}, 'box_15_2_employee_state_id': {'title': 'Box 15 2 Employee State Id', 'type': 'string'}, 'box_16_2_state_wages': {'title': 'Box 16 2 State Wages', 'type': 'number'}, 'box_17_2_state_income_tax': {'title': 'Box 17 2 State Income Tax', 'type': 'number'}, 'box_18_2_local_wages': {'title': 'Box 18 2 Local Wages', 'type': 'number'}, 'box_19_2_local_income_tax': {'title': 'Box 19 2 Local Income Tax', 'type': 'number'}, 'box_20_2_locality': {'title': 'Box 20 2 Locality', 'type': 'string'}}, 'required': ['box_b_employer_identification_number', 'box_c_employer_name', 'box_c_employer_street_address', 'box_c_employer_city_state_zip', 'box_a_employee_ssn', 'box_e_employee_name', 'box_e_employee_street_address', 'box_e_employee_city_state_zip', 'box_d_control_number', 'box_1_wages', 'box_2_federal_tax_withheld', 'box_3_social_security_wages', 'box_4_social_security_tax_withheld', 'box_5_medicare_wages', 'box_6_medicare_wages_tax_withheld', 'box_7_social_security_tips', 'box_8_allocated_tips', 'box_9_advance_eic_payment', 'box_10_dependent_care_benefits', 'box_11_nonqualified_plans', 'box_12a_code', 'box_12a_value', 'box_12b_code', 'box_12b_value', 'box_12c_code', 'box_12c_value', 'box_12d_code', 'box_12d_value', 'box_13_statutary_employee', 'box_13_retirement_plan', 'box_13_third_part_sick_pay', 'box_15_1_state', 'box_15_1_employee_state_id', 'box_16_1_state_wages', 'box_17_1_state_income_tax', 'box_18_1_local_wages', 'box_19_1_local_income_tax', 'box_20_1_locality', 'box_15_2_state', 'box_15_2_employee_state_id', 'box_16_2_state_wages', 'box_17_2_state_income_tax', 'box_18_2_local_wages', 'box_19_2_local_income_tax', 'box_20_2_locality'], 'title': 'W2Form', 'type': 'object'}"
6+
sample_percent: 1 # % of the dataset to use; 1.0 means use the entire dataset
7+
resplit_train_percent: 0.3 # % of the sampled dataset to use for training; the rest is used for validation
8+
image_column: image
9+
user_text_column: null
10+
assistant_text_column: ground_truth
11+
grader: JSONGrader # Task-specific grader
12+
13+
14+
task2:
15+
dataset: getomni-ai/ocr-benchmark
16+
is_local: false
17+
system_prompt: You are a helpful assistant, you will always respond only in JSON following the provided JSON schema.
18+
user_prompt: "Extract the data in this image as a JSON. Use the following JSON schema:\n"
19+
sample_percent: 0.6
20+
resplit_train_percent: 0.0
21+
image_column: image
22+
user_text_column: json_schema
23+
assistant_text_column: true_json_output
24+
grader: JSONGrader
25+
26+
27+
finetuning:
28+
### FFT LAYERS TO TRAIN - ALL FALSE FOR NO FFT
29+
fusion: true
30+
fusion+encoder: false
31+
fusion+decoder: false
32+
fusion+encoder+decoder: true
33+
### LORA RANKS TO TRAIN - EMPTY LIST FOR NO LORA
34+
lora_ranks: [8, 64]
35+
### TORCHTUNE CONFIG
36+
fft_torchtune_config: transferability/finetune/8b_full.yaml
37+
lora_torchtune_config: transferability/finetune/8b_lora.yaml
38+
### TORCHTUNE ARGS
39+
model_path: /path/to/llama31/ckpt
40+
tokenizer_path: /path/to/llama31/ckpt/tokenizer.model
41+
epochs: 5 # Number of training epochs
42+
batch_size: 8 # Batch size per device for training
43+
ngpu: 4
44+
distributed: true # Whether to use distributed training
45+
46+
47+
48+
evals:
49+
nb_eval_samples: null # Number of samples to use for evaluation; null means use the entire dataset.
50+
checkpoint_to_eval: -1
51+
model_server_args:
52+
tensor_parallel_size: 2
53+
max_model_len: 8192
54+
max_num_seqs: 128
55+
enforce_eager: true
56+
inference_params:
57+
temperature: 0
58+
top_p: 1.0
59+
max_completion_tokens: 4096
60+
seed: 42

0 commit comments

Comments
 (0)