wehub-resource-sync
159 lines
4.9 KiB
Python
159 lines
4.9 KiB
Python
import os
|
|
import json
|
|
import logging
|
|
from dataclasses import dataclass, field
|
|
from typing import Optional
|
|
|
|
import torch
|
|
from fairseq import distributed_utils, utils
|
|
from fairseq.dataclass import FairseqDataclass
|
|
from fairseq.models import (
|
|
FairseqIncrementalDecoder,
|
|
FairseqLanguageModel,
|
|
register_model,
|
|
register_model_architecture,
|
|
)
|
|
|
|
from omegaconf import II
|
|
|
|
from fairseq.model_parallel.megatron.mpu import (
|
|
initialize_model_parallel,
|
|
model_parallel_is_initialized
|
|
)
|
|
from .decoder.yoco import YOCO, YOCOArgs
|
|
|
|
DEFAULT_MAX_TARGET_POSITIONS = 4096
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class LanguageConfig(FairseqDataclass):
|
|
yoco_model: Optional[str] = field(
|
|
default=None,
|
|
metadata={"help": "path to load params from"},
|
|
)
|
|
load_ckpt: Optional[str] = field(
|
|
default=None,
|
|
metadata={"help": "path to load checkpoint from"},
|
|
)
|
|
dim: int = field(
|
|
default=1024,
|
|
)
|
|
hidden_dim: int = field(
|
|
default=3072,
|
|
)
|
|
n_layers: int = field(
|
|
default=24,
|
|
)
|
|
n_self_heads: int = field(
|
|
default=4,
|
|
)
|
|
n_attn_heads: int = field(
|
|
default=8,
|
|
)
|
|
n_attn_kv_heads: Optional[int] = field(
|
|
default=None,
|
|
)
|
|
batch_size: int = field(
|
|
default=1,
|
|
)
|
|
share_input_output_embed: bool = field(
|
|
default=False, metadata={"help": "share decoder input and output embeddings"}
|
|
)
|
|
sliding_window: Optional[bool] = field(
|
|
default=None,
|
|
)
|
|
rope_theta: Optional[float] = field(
|
|
default=10000.0,
|
|
)
|
|
checkpoint_activations: bool = field(
|
|
default=False, metadata={"help": "checkpoint activations at each layer"}
|
|
)
|
|
tokens_per_sample: int = II("task.tokens_per_sample")
|
|
model_parallel_size: int = II("common.model_parallel_size")
|
|
|
|
|
|
@register_model("yoco", dataclass=LanguageConfig)
|
|
class LanguageModel(FairseqLanguageModel):
|
|
def __init__(self, args, decoder, tokenizer):
|
|
self.args = args
|
|
self.tokenizer = tokenizer
|
|
super().__init__(decoder)
|
|
|
|
@classmethod
|
|
def build_model(cls, args, task):
|
|
if not model_parallel_is_initialized():
|
|
initialize_model_parallel(args.model_parallel_size)
|
|
|
|
if args.yoco_model is None:
|
|
params = {
|
|
"dim": args.dim,
|
|
"n_layers": args.n_layers,
|
|
"n_self_heads": args.n_self_heads,
|
|
"n_attn_heads": args.n_attn_heads,
|
|
"n_attn_kv_heads": args.n_attn_kv_heads,
|
|
"hidden_dim": args.hidden_dim,
|
|
"vocab_size": task.tokenizer.n_words,
|
|
"max_batch_size": args.batch_size,
|
|
"max_seq_len": args.tokens_per_sample,
|
|
"model_parallel_size": args.model_parallel_size,
|
|
"load_checkpoint": args.load_ckpt is not None,
|
|
"rope_theta": args.rope_theta,
|
|
}
|
|
model_args: YOCOArgs = YOCOArgs(
|
|
**params,
|
|
)
|
|
else:
|
|
with open(os.path.join(args.yoco_model, "params.json"), "r") as f:
|
|
params = json.load(f)
|
|
model_args = YOCOArgs(**params)
|
|
model_args.max_batch_size = args.batch_size
|
|
model_args.max_seq_len = args.tokens_per_sample
|
|
model_args.model_parallel_size = args.model_parallel_size
|
|
model_args.load_checkpoint = args.load_ckpt is not None
|
|
|
|
model = YOCO(
|
|
model_args,
|
|
checkpoint_activations=args.checkpoint_activations,
|
|
)
|
|
if args.load_ckpt is not None:
|
|
loaded = torch.load(args.load_ckpt, mmap=True)
|
|
model.load_state_dict(loaded, assign=True)
|
|
model = YOCOModel(model)
|
|
return cls(args, model, task.tokenizer)
|
|
|
|
class YOCOModel(FairseqIncrementalDecoder):
|
|
def __init__(self, model):
|
|
super().__init__(None)
|
|
self.model = model
|
|
|
|
def forward(self, src_tokens, **kwargs):
|
|
return self.model.forward(src_tokens, **kwargs)
|
|
|
|
def max_positions(self):
|
|
return self.model.args.max_seq_len
|
|
|
|
def default(args):
|
|
args.n_attn_kv_heads = getattr(args, "n_attn_kv_heads", args.n_attn_heads)
|
|
args.sliding_window = getattr(args, "sliding_window", False)
|
|
args.rope_theta = getattr(args, "rope_theta", 10000.0)
|
|
args.share_input_output_embed = getattr(
|
|
args, "share_input_output_embed", False
|
|
)
|
|
args.checkpoint_activations = getattr(args, "checkpoint_activations", False)
|
|
|
|
|
|
@register_model_architecture("yoco", "yoco_3b")
|
|
def yoco_3b(args):
|
|
args.dim = getattr(args, "dim", 3072)
|
|
args.hidden_dim = getattr(args, "hidden_dim", 8192)
|
|
args.n_layers = getattr(args, "n_layers", 26)
|
|
args.n_self_heads = getattr(args, "n_self_heads", 24)
|
|
args.n_attn_heads = getattr(args, "n_attn_heads", 24)
|
|
args.n_attn_kv_heads = getattr(args, "n_attn_kv_heads", 8)
|
|
default(args)
|
|
|
|
|
|
|
|
|