| |
| import argparse |
| import json |
| import os |
| import sys |
|
|
| import torch |
| from transformers import GPT2LMHeadModel |
|
|
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| from model import GPT, GPTConfig |
|
|
|
|
| TRANSPOSED = ( |
| 'attn.c_attn.weight', |
| 'attn.c_proj.weight', |
| 'mlp.c_fc.weight', |
| 'mlp.c_proj.weight', |
| ) |
|
|
|
|
| def copy_tensor(dst, src, name): |
| value = src.t() if any(name.endswith(suffix) for suffix in TRANSPOSED) else src |
| if dst.shape != value.shape: |
| raise ValueError(f'shape mismatch for {name}: dst={tuple(dst.shape)} src={tuple(value.shape)}') |
| dst.copy_(value) |
|
|
|
|
| def copy_hf_block(model, hf_state, src_layer, dst_layer): |
| own_state = model.state_dict() |
| prefix_hf = f'transformer.h.{src_layer}.' |
| prefix_own = f'transformer.h.{dst_layer}.' |
| keys = [ |
| 'ln_1.weight', |
| 'ln_1.bias', |
| 'attn.c_attn.weight', |
| 'attn.c_attn.bias', |
| 'attn.c_proj.weight', |
| 'attn.c_proj.bias', |
| 'ln_2.weight', |
| 'ln_2.bias', |
| 'mlp.c_fc.weight', |
| 'mlp.c_fc.bias', |
| 'mlp.c_proj.weight', |
| 'mlp.c_proj.bias', |
| ] |
| for suffix in keys: |
| own_key = prefix_own + suffix |
| hf_key = prefix_hf + suffix |
| copy_tensor(own_state[own_key], hf_state[hf_key], own_key) |
|
|
|
|
| def zero_residual_outputs(model, layer_ids): |
| with torch.no_grad(): |
| for i in layer_ids: |
| block = model.transformer.h[i] |
| block.attn.c_proj.weight.zero_() |
| block.mlp.c_proj.weight.zero_() |
| if block.attn.c_proj.bias is not None: |
| block.attn.c_proj.bias.zero_() |
| if block.mlp.c_proj.bias is not None: |
| block.mlp.c_proj.bias.zero_() |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--out_dir', default='out-init-distilgpt2-plus2') |
| parser.add_argument('--model_name', default='distilgpt2') |
| parser.add_argument('--dropout', type=float, default=0.0) |
| args = parser.parse_args() |
|
|
| os.makedirs(args.out_dir, exist_ok=True) |
|
|
| hf = GPT2LMHeadModel.from_pretrained(args.model_name) |
| hf.eval() |
| hf_state = hf.state_dict() |
| hf_cfg = hf.config |
| if hf_cfg.vocab_size != 50257 or hf_cfg.n_layer != 6 or hf_cfg.n_head != 12 or hf_cfg.n_embd != 768: |
| raise ValueError( |
| 'Expected distilgpt2 shape 6L/12H/768d/vocab50257, got ' |
| f'{hf_cfg.n_layer}L/{hf_cfg.n_head}H/{hf_cfg.n_embd}d/vocab{hf_cfg.vocab_size}' |
| ) |
|
|
| model_args = dict( |
| block_size=1024, |
| vocab_size=50257, |
| n_layer=8, |
| n_head=12, |
| n_embd=768, |
| dropout=args.dropout, |
| bias=True, |
| ) |
| torch.manual_seed(1337) |
| model = GPT(GPTConfig(**model_args)) |
|
|
| with torch.no_grad(): |
| own_state = model.state_dict() |
| own_state['transformer.wte.weight'].copy_(hf_state['transformer.wte.weight']) |
| own_state['transformer.wpe.weight'].copy_(hf_state['transformer.wpe.weight']) |
| own_state['transformer.ln_f.weight'].copy_(hf_state['transformer.ln_f.weight']) |
| own_state['transformer.ln_f.bias'].copy_(hf_state['transformer.ln_f.bias']) |
| own_state['lm_head.weight'].copy_(hf_state['transformer.wte.weight']) |
| for i in range(6): |
| copy_hf_block(model, hf_state, i, i) |
| zero_residual_outputs(model, [6, 7]) |
|
|
| total_params = sum(p.numel() for p in model.parameters()) |
| if total_params >= 100_000_000: |
| raise ValueError(f'model has {total_params:,} params, exceeds 100M limit') |
|
|
| checkpoint = { |
| 'model': model.state_dict(), |
| 'optimizer': None, |
| 'model_args': model_args, |
| 'iter_num': 0, |
| 'best_val_loss': 1e9, |
| 'config': { |
| 'source_model': args.model_name, |
| 'architecture_note': 'distilgpt2 6 pretrained layers + 2 zero-output residual layers', |
| 'total_params': total_params, |
| }, |
| } |
| torch.save(checkpoint, os.path.join(args.out_dir, 'checkpoint.pt')) |
| with open(os.path.join(args.out_dir, 'config.json'), 'w', encoding='utf-8') as f: |
| json.dump(model_args, f, indent=2) |
| with open(os.path.join(args.out_dir, 'init_info.json'), 'w', encoding='utf-8') as f: |
| json.dump(checkpoint['config'], f, indent=2) |
| print(f'wrote {args.out_dir}/checkpoint.pt') |
| print(f'total params: {total_params:,}') |
| print('tokenizer/vocab: GPT-2 ids, vocab_size=50257') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
