code/reproduce_hns_seed_checkpoint.py

#!/usr/bin/env python3
"""Portable helper shipped with each published LoRA; matches effective zero warmup."""
import argparse
import hashlib
import json
import os
import subprocess
import sys
import tempfile
from pathlib import Path


def sha256(path):
    h = hashlib.sha256()
    with path.open('rb') as f:
        for chunk in iter(lambda: f.read(8 * 1024 * 1024), b''):
            h.update(chunk)
    return h.hexdigest()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('action', choices=['train', 'build-hns', 'evaluate'])
    parser.add_argument('--data-dir', type=Path)
    parser.add_argument('--grid-dir', type=Path)
    parser.add_argument('--output-dir', required=True, type=Path)
    parser.add_argument('--base-dir', type=Path)
    parser.add_argument('--peer-root', type=Path)
    parser.add_argument('--token-budget', type=int)
    args = parser.parse_args()
    repo = Path(__file__).resolve().parents[1]
    if args.output_dir.resolve().is_relative_to(repo):
        parser.error('--output-dir must be outside the downloaded repository to avoid recursive grid copies')
    meta = json.loads((repo / 'publication.json').read_text())
    source = repo / 'source_snapshot'
    if not source.is_dir():
        raise RuntimeError('Extract code/source_snapshot.tar.gz in the repository directory first.')
    env = dict(os.environ)
    env['PYTHONPATH'] = str(source / 'src') + os.pathsep + str(source / 'scripts')
    env.update(TOKENIZERS_PARALLELISM='false', VLLM_BATCH_INVARIANT='1', VLLM_USE_FLASHINFER_SAMPLER='0',
               CUBLAS_WORKSPACE_CONFIG=':4096:8', PYTHONHASHSEED=str(meta['training_seed'] if args.action == 'train' else 42),
               VLLM_DISABLE_COMPILE_CACHE='1')
    with tempfile.TemporaryDirectory(prefix='hr-') as temporary:
        env.update(TMPDIR=temporary, TMP=temporary, TEMP=temporary, VLLM_CACHE_ROOT=temporary + '/vc')
        base = args.base_dir
        if args.action != 'build-hns' and base is None:
            from huggingface_hub import snapshot_download
            revision = meta['base_provenance']['revision']
            if revision is None:
                raise RuntimeError('A unique base revision is unavailable; provide --base-dir and verify publication.json hashes.')
            base = Path(snapshot_download(meta['base_provenance']['model_id'], revision=revision))
        if base is not None:
            for name, expected in meta['base_provenance']['files_sha256'].items():
                assert sha256(base / name) == expected, f'Base file mismatch: {name}'
        task = meta['train_task']
        if args.action == 'train':
            if args.data_dir is None:
                parser.error('train requires --data-dir')
            dataset = args.data_dir / meta['dataset']['filename']
            assert sha256(dataset) == meta['dataset']['file_sha256'], 'Training dataset hash mismatch'
            a = meta['requested_run_args']
            command = [sys.executable, '-m', 'finetune.train_sft_peft', '--base_model', str(base), '--task', 'if' if task == 'tulu' else task,
                       '--peft_method', 'lora', '--dataset_path', str(dataset), '--output_dir', str(args.output_dir), '--train_profile', 'none',
                       '--warmup_ratio', '0', '--seed', str(meta['training_seed']), '--dataset_seed', '42', '--bf16', '--gradient_checkpointing']
            fields = ['num_train_epochs', 'per_device_train_batch_size', 'gradient_accumulation_steps', 'global_train_batch_size', 'lr', 'weight_decay',
                      'min_lr_ratio', 'lr_scheduler_type', 'adam_beta1', 'adam_beta2', 'grad_clip', 'max_seq_len', 'max_train_samples', 'sft_format',
                      'chat_template_mode', 'target_modules', 'r', 'lora_alpha', 'lora_dropout', 'pad_to_multiple_of', 'save_strategy', 'save_steps',
                      'save_total_limit', 'logging_steps']
            for key in fields:
                if a.get(key) is not None:
                    command.extend(['--' + key, str(a[key])])
            subprocess.run(command, env=env, check=True)
        else:
            benchmark_meta = json.loads((repo / 'evaluation/benchmark_provenance.json').read_text())
            benchmark = repo / 'evaluation/benchmark_input' / benchmark_meta['filename']
            assert sha256(benchmark) == benchmark_meta['sha256']
            cfg = {'bases': {meta['base']: str(base)}, 'checkpoints': [{'base': meta['base'], 'train_task': task, 'lora': str(repo)}],
                   'evaluation_tasks': {task: {**meta['evaluation_task'], 'dataset': str(benchmark)}}}
            if args.action == 'build-hns':
                cfg['checkpoints'] = []
                for peer in meta['peer_source_adapters']:
                    if peer['task'] == task:
                        peer_dir = repo
                    elif args.peer_root is not None:
                        peer_dir = args.peer_root / peer['repo_id'].split('/')[-1]
                    else:
                        from huggingface_hub import snapshot_download
                        peer_dir = Path(snapshot_download(peer['repo_id'], allow_patterns=['adapter*', 'tokenizer*', 'chat_template*', 'special_tokens_map.json', 'publication.json']))
                    assert sha256(peer_dir / 'adapter_model.safetensors') == peer['weights_sha256'], f'Peer adapter mismatch: {peer["repo_id"]}'
                    cfg['checkpoints'].append({'base': meta['base'], 'train_task': peer['task'], 'lora': str(peer_dir)})
            config = Path(temporary) / 'task.json'
            config.write_text(json.dumps(cfg))
            if args.action == 'build-hns':
                command = [sys.executable, str(source / 'scripts/build_hns_step_grid_2x4.py'), '--task_config', str(config),
                           '--grid_config', str(repo / 'hns/grid_config.json'), '--base', meta['base'], '--output_dir', str(args.output_dir), '--device', 'cuda']
                subprocess.run(command, env=env, check=True)
            else:
                if args.grid_dir is None:
                    parser.error('evaluate requires --grid-dir')
                manifest = args.grid_dir / meta['base'] / 'variant_manifest.json'
                for script, flags in [
                    ('eval_forgetting_matrix_vllm.py', ['--base_model', str(base), '--variant_manifest', str(manifest), '--config', str(config),
                     '--output_dir', str(args.output_dir), '--tasks', task, '--diagonal_only', '--max_model_len', '4096', '--gpu_memory_utilization', '0.94',
                     '--max_num_seqs', '1024', '--max_num_batched_tokens', str(args.token_budget or meta['inference']['max_num_batched_tokens']),
                     '--adapter_block_size', '11', '--max_lora_rank', '16', '--prompt_chunk_short', '4096', '--prompt_chunk_long', '512', '--seed', '42']),
                    ('score_forgetting_matrix.py', ['--matrix_dir', str(args.output_dir), '--workers', '32']),
                ]:
                    subprocess.run([sys.executable, str(source / 'scripts' / script), *flags], env=env, check=True)


if __name__ == '__main__':
    main()