How to train in mult-node?

[Nightbringers]

I dont konw how to train use mult-host. Can you gave a example? thank you.

[erfanzar]

hello
yes for sure but which backend you want to use for multi host training is it on multiple TPU pods or GPU servers?

[Nightbringers]

GPU servers and slurm cluster

[erfanzar]

can you please test this code

import jax
from EasyDel import TrainArguments, CausalLMTrainer

num_processes = 6
process_id = 0  # number between 0 and num_processes-1 that says which node is current node
coordinator_address = 'ip:port'  # for example 192.168.1.12:8600 (make sure this port is not closed by firewall)
jax.distributed.initialize(coordinator_address=coordinator_address,
                           num_processes=num_processes,
                           process_id=process_id)

train_args = TrainArguments(
    backend='gpu',
    sharding_array=(num_processes, -1, 1),
    use_wandb=True,
    use_pjit_attention_force=False
)

trainer = CausalLMTrainer(
    arguments=train_args,
    dataset_train=..., # To Be passed
    ckpt_path=... # To Be passed path to ckpt or None 
)
parameters = None  # if you want to finetune a model you can pass parameters to trainer and they should be like frozen({"params":...})
trainer.train(model_parameters=parameters or None)

[Nightbringers]

I encountered some issues when I tried to run EasyDeL/examples/training/causal-lm/llama.py. I use llama-13b. I first convert hf-llama to flax. Use this code:

model = AutoModelForCausalLM.from_pretrained(path)
state_dict = model.state_dict()
flax_params = llama_convert_hf_to_flax(state_dict, num_hidden_layers=40, num_attention_heads=40, hidden_size=5120,device = device)
save_ckpt(flax_params, 'flax_param_easydel')

Is my convert code correct?

then I run llama.py, here is the error:
ValueError: One of pjit outputs was given the sharding of NamedSharding(mesh={'dp': 1, 'fsdp': 8, 'mp': 1}, spec=PartitionSpec('fsdp',)), which implies that the global size of its dimension 0 should be divisible by 8, but it is equal to 40076 (full shape: (40076, 5120))

[erfanzar]

when you trying to run pass fully_fsdp=False in config.get_partition_rules that will fix this problem

[Nightbringers]

This seems to disable fully_fsdp. But I want use fully_fsdp.

[erfanzar]

which model you trying to use
can you give me the config of model that you trying to use so i can write a custom partition rule for that

[Nightbringers]

llama-13b v2. But the model's vocabulary size has changed. Does this have any impact?

[erfanzar]

yes that's the reason that you get this error can you tell me your vocab size? like is it like EOS,BOS added version that have 32002 tokens?

[erfanzar]

use this

from jax.sharding import PartitionSpec as PS
from EasyDel import TrainArguments

partition_rules = (

    ("transformer/wte/embedding", PS('dp', "fsdp")),

    ("attention/(wq|wk|wv)/kernel", PS("fsdp")),
    ("attention/wo/kernel", PS("fsdp")),

    ("feed_forward/w1/kernel", PS("fsdp")),
    ("feed_forward/w2/kernel", PS("fsdp")),
    ("feed_forward/w3/kernel", PS("fsdp")),

    ("attention_norm/kernel", PS('fsdp')),
    ("ffn_norm/kernel", PS('fsdp')),

    ("transformer/ln_f/kernel", PS('fsdp')),
    ("lm_head/kernel", PS("fsdp", 'dp')),
    ('.*', PS('fsdp')),
)

train_args = TrainArguments(
    custom_rule=partition_rules,
    ...
)

this one have to work fine if you just have changed the vocab size

[Nightbringers]

yes, it worked!

[erfanzar]

if you have any other issue please let me know <3

[Nightbringers]

The calculation of losses has encountered an issue. The error is:
File "/EasyDeL/EasyDel/trainer/fsdp_train.py", line 406, in train
sharded_train_state_, loss, accuracy = self.sharded_train_step_fn(sharded_train_state_,
File "/EasyDeL/EasyDel/trainer/fsdp_train.py", line 309, in fsdp_train_step_
(loss__, accuracy__), grad = grad_fn(state.params)
File "/EasyDeL/EasyDel/trainer/fsdp_train.py", line 303, in calculate_loss
loss, accuracy = loss_fn(
File "/anaconda3-2/envs/py3.10/lib/python3.10/site-packages/fjutils/easylm.py", line 495, in blockwise_cross_entropy
logits = rearrange(logits, '(n c) d -> n c d', c=chunk_size)
File "/anaconda3-2/envs/py3.10/lib/python3.10/site-packages/einops/einops.py", line 483, in rearrange
return reduce(cast(Tensor, tensor), pattern, reduction='rearrange', **axes_lengths)
File "/anaconda3-2/envs/py3.10/lib/python3.10/site-packages/einops/einops.py", line 420, in reduce
raise EinopsError(message + '\n {}'.format(e))
einops.EinopsError: Error while processing rearrange-reduction pattern "(n c) d -> n c d".
Input tensor shape: (16368, 55296). Additional info: {'c': 1024}.
Shape mismatch, can't divide axis of length 16368 in chunks of 1024

my vocab size is 55296, sequence_length is 1024.

[erfanzar]

set loss_remat to ''

train_args = TrainArguments(
    custom_rule=partition_rules,
    loss_remat=''
)

this will work, the current error that you takin is because you trying to use blockwise crossentropy and your vocab size (55296) is not visible by 1024 so you can either change loss_remat to '' or change your loss_chunk

[Nightbringers]

    if self.arguments.loss_remat != '':
        blockwise_cross = functools.partial(
            blockwise_cross_entropy,
            chunk_size=self.arguments.loss_chunk,
            policy=self.arguments.loss_remat
        )
        loss_fn = blockwise_cross
    else:
        loss_fn = cross_entropy_loss_and_accuracy

What is the difference between blockwise_cross and cross_entropy_loss_and_accuracy? What is the difference between blockwise crossentropy and crossentropy? Are there any advantages to using blockwise_cross？