Open xrsrke opened 9 months ago
I'm taking a stab at this issue. I'll first come up with a solution that works for at least bloom-560m
, which is already mapped in ParallelMapping
, such that we can see if the automatic mapping works. Then we can work on making it more general/generalizable.
I'm working on this issue.
For testing, can I see some MWEs of what the input of these methods looks like? Like, typical inputs to make sure I get the right results.
@yugen-ok
model = AutoModel.from_pretrained()
parallel_mapping = ParallelMapping(model)
ref_mapping = {...}
for name, _ in model.named_modules():
mapping[name] = parallel_mapping.is_row_parallel(name, module)
assert ref_mapping == mapping
Notes
pipegoose.nn.parallel_mapping.ParallelMapping
module
is an instance inmodel.named_modules()
AutoModelForCausalLM.from_pretrained()
,torch.nn.Transformer
APIs
Write a function that:
ParallelMapping(model).is_column_parallel(name, module)
returns True if the module is the first linear layer in an MLP layer, or if the module is a query, key, value linear, or a fused qkv linear of an attention layer, or an input embedding; otherwise, return False.ParallelMapping(model).is_row_parallel(name, module)
returns True if the module is the second linear layer in an MLP layer, or the output projection of an attention layer.ParallelMapping(model).is_lm_head(name, module)
returns True if the module is language model head.ParallelMapping(model).is_text_embedding(name, module)
returns True if the module is text embedding module.ParallelMapping(model).is_mlp(name, module)
returns True if the module is an MLP layer