Llama 400B en agustina

He estado preparando una demo para correr el fp16 del Llama3.1 grande en nuestro fragmento de tarjetas L40s. Es basicamente el ejemplo que aparece en https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/llama con algunos detallitos que es facil saltarse.

Lo primero es decidir la paralelizacion. Usando 24 GPUs me ha dado ocasionalmente problemas de vRAM asi que por ir a lo seguro conviene llenar la particion -dejando una maquina libre- y por tanto usar 8 maquinas con 4 GPUs cada una. Esto nos permite hacer Tensor Parallelism = 4 y Pipeline Parallelism = 8. Y dejamos el paralelismo de contexto a 1, porque en el ejemplo tampoco lo mencionan mucho.

Lo primero es pues preparar el modelo siguiendo las instrucciones pero usando una de nuestras tarjetas graficas. Para ello primero se prepara el checkpoint. Aqui usamos srun + mpirun, lanzando por tanto en interactivo… hay que tener cuidado que no se corte la terminal, porque las tareas llevan horas.

cd TensorRT-LLM/examples/llama/
ln -s  ../../../.cache/huggingface/hub/models--meta-llama--Llama-3.1-405B-Instruct/snapshots/be673f326cab4cd22ccfef76109faf68e41aa5f1/ llama400 
srun --account=... --partition=...  -N 1 -n 1 --nodelist=...  --gpus=4  --time=05:00:00 mpirun -n 1 python convert_checkpoint.py --dtype float16 --tp_size 4 --pp_size 8  --model_dir llama400 --output_dir llama400_8_4.ckpt

y luego propiamente el modelo. Ojo, el modelo se crea con la misma arquitectura que las tarjetas en la que se ejecute la compilacion, y las L40s son la arquitectura anterior a las H100, luego hay que ejecutarlo en una particion de las L40s

srun --account=... --partition=ada -N 1 --ntasks-per-node=1 --gpus=4 --time=12:00:00 mpirun trtllm-build --checkpoint_dir llama400_8_4.ckpt --output_dir llama400_8_4.adaengine

Despues de esto viene el gran truco: las instrucciones por defecto crean el modelo asumiendo 8 GPUs por maquina y el pytorch se queja. Esto se debe aque si bien el build tiene el comando –gpus_per_node, el mapping ya lo ha creado el conversor, que no lo tiene previsto. Asi que hay que editar bien el config.json del engine bien el del checkpoint.

        "mapping": {
            "world_size": 32,
            "gpus_per_node": 4,
            "cp_size": 1,
            "tp_size": 4,
            "pp_size": 8,
            "moe_tp_size": 4,
            "moe_ep_size": 1
        },

y ahora a ejecutar.

NO HE CONSEGUIDO EJECUTARLO VIA SBATCH. Tengo que asignar primero los recursos con salloc y lanzar tras haber configurado correctamente los modulos (solo hay soporte para python 3.10 en el codigo liberado por nvidia). Basicamente

salloc --account=... --partition=ada -N 8 --ntasks-per-node=1 --gpus=32 --time=23:00:00

Y ya dentro de la sesion interactiva arreglar los modulos y lanzar:

module use /opt/amd/spack/share/spack/modules/linux-rocky8-zen3
module load openmpi/4.1.5
module load python
module load nvidia-hpc-sdk/24.5
module load numactl/2.0.14
cd TensorRT-LLM/examples/llama/
mpirun -v --display-allocation --display-map --report-pid - --report-uri - -n 32 --oversubscribe python ../run.py --engine_dir llama400_8_4.adaengine --tokenizer_dir llama400  --max_output_len 100  --input_text "A que instituto de Unizar corresponden las siglas BIFI?"

El mpirun cumple y lanza todo sin sorpresas.