universal-ckp: support megatron-deepspeed llama model (#4666)
Megatron-DeepSpeed's llama implementation of swiglu allocates a single
ColumnParallelLinear layer L, but effectively this parameter is a
container of two Linear layers L1, L2 used for silu(L1(x)) * L2(x)).
This requires special handling in ds_to_universal to create a
representation of L parameter where the slices of L1 and L2 are first
concatenated and then L is created by concatenating L1 and L2.
Signed-off-by: Moshe Island <misland@habana.ai>
Co-authored-by: Moshe Island <misland@habana.ai>
Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
