Пример использования PyTorch DDP

Разберем предоставленный код подробно, объясняя каждый блок и его роль в процессе распределенного обучения с использованием PyTorch Distributed Data Parallel (DDP).

Примечание

Полный исходный код доступен для копирования в репозитории на GitHub.

Код можно разделить на несколько основных частей:

  • настройка среды;

  • инициализация процесса обучения;

  • основная функция обучения;

  • завершение работы.

Основные понятия

  • Node Rank (ранг узла) — идентификатор узла в кластере.

    В контексте DDP каждый узел обработки (например, каждый GPU) имеет свой уникальный ранг в глобальном пространстве всех узлов. Ранг можно получить с помощью dist.get_rank().

  • World Size (размер мира) — общее количество узлов или процессов в распределенной системе.

    Это число указывает на общее количество процессов, участвующих в обучении. Получить размер мира можно с помощью dist.get_world_size().

  • Local Rank (локальный ранг) — идентификатор узла в пределах одного физического сервера или машины.

    Если вы используете несколько GPU на одном сервере, каждый GPU будет иметь свой локальный ранг. В примере кода LOCAL_RANK получается из переменной окружения LOCAL_RANK, что является стандартным способом его определения при использовании скриптов запуска, предоставляемых PyTorch, таких как torch.distributed.launch или torchrun.

Шаг 1. Настройка среды

Рабочая директория определяется как абсолютный путь до каталога, в котором находится исполняемый файл. Это упрощает доступ к файлам и директориям относительно местоположения скрипта.

Логирование настраивается для записи в файл train.log событий процесса обучения: времени, уровня логирования и сообщений, что помогает в отладке и мониторинге процесса обучения.

Установка рабочей директории и логирование

import   os 
 import   torch 
 from   torch.nn.parallel   import   DistributedDataParallel   as   DDP 
 import   torch.distributed   as   dist 
 import   pathlib 
 import   logging 
 from   datasets   import   load_dataset 
 from   peft   import   ( 
     LoraConfig , 
     get_peft_model , 
     prepare_model_for_kbit_training 
 ) 
 from   transformers   import   ( 
     AutoModelForCausalLM , 
     AutoTokenizer , 
     BitsAndBytesConfig , 
     TrainingArguments , 
     Trainer , 
     DataCollatorForLanguageModeling 
 ) 


 BASE_DIR   =   str ( pathlib . Path ( __file__ ) . parent . absolute ()) 


 print ( f "Working dir:  { BASE_DIR } " ) 






 logging . basicConfig ( filename = BASE_DIR   +   "/train.log" , 


                     filemode = 'a' , 


                     format = ' %(asctime)s , %(msecs)d   %(name)s   %(levelname)s   %(message)s ' , 


                     datefmt = '%H:%M:%S' , 


                     level = logging . INFO ) 


 logging . info ( f "Working dir:  { BASE_DIR } " ) 


 def   cleanup (): 
     dist . destroy_process_group ()

Шаг 2. Инициализация распределенного обучения

Переменные окружения WORLD_SIZE и WORLD_RANK получаются из переменных окружения, заданных внешне, и определяют общее количество процессов и ранг текущего процесса соответственно.

Инициализация группы процессов init_processes создает группу процессов для распределенного обучения, используя выбранный бэкенд (по умолчанию — nccl для NVIDIA GPUs). Затем вызывает функцию обучения с текущим рангом процесса, общим количеством процессов и локальным рангом.

def   init_processes ( fn ,   local_rank ,   backend = 'nccl' ): 


     dist . init_process_group ( backend ) 


     fn ( dist . get_rank (),   dist . get_world_size (),   local_rank )

Шаг 3. Запуск основной функции обучения

Включает в себя следующие шаги:

  • Настройка модели и токенизатора.

    Загружаются модель и токенизатор, настраивается квантизация модели для уменьшения занимаемой памяти и увеличения скорости вычислений.

  • Распределенное обучение с DDP.

    Модель оборачивается в DDP, что позволяет синхронизировать градиенты между процессами. device_ids=[local_rank] указывает, что каждый процесс использует свой собственный GPU.

  • Подготовка данных.

    Загрузка и токенизация данных для обучения.

  • Настройка аргументов обучения и тренера.

    Создание объекта Trainer с заданными параметрами обучения для управления процессом обучения.

  • Сохранение модели.

    После обучения состояние модели сохраняется на диск. Это позволяет использовать модель после обучения.

Важно понимать роль node rank, world size и local rank для корректной работы в распределенной системе. Приведенный код демонстрирует основные шаги настройки и использования DDP для обучения модели языка с использованием технологий Hugging Face и PyTorch.

def   train ( rank ,   size ,   local_rank ): 


     epochs   =   10    # Number of training epochs 
     MODEL_NAME   =   "Intel/neural-chat-7b-v3-1" 
     # Configuration for quantizing the model 
     bnb_config   =   BitsAndBytesConfig ( 
         load_in_4bit = True ,    # Download in 4-bit format 
         bnb_4bit_use_double_quant = True ,    # Using double quantization 
         bnb_4bit_quant_type = "nf4" ,    # Quantization type 
         bnb_4bit_compute_dtype = torch . bfloat16    # Calculation data type 
     ) 


     # Loading a pre-trained model with quantization settings 
     model   =   AutoModelForCausalLM . from_pretrained ( 
         MODEL_NAME , 
         device_map = "cuda" ,    # Distribution of the model by CUDA devices 
         trust_remote_code = True ,    # Trusting remote code 
         quantization_config = bnb_config    # Applying quantize configuration 
     ) 


     # Loading a tokenizer for a model 
     tokenizer   =   AutoTokenizer . from_pretrained ( MODEL_NAME ) 
     tokenizer . pad_token   =   tokenizer . eos_token    # Setting the padding token equal to the end of line token 


     model . gradient_checkpointing_enable ()    # Enable gradient control points to save memory 
     model   =   prepare_model_for_kbit_training ( model )    # Preparing the model for training with kbit optimization 


     # Configuration for optimization LoRA 
     config   =   LoraConfig ( 
         r = 8 ,    # Projection dimension 
         lora_alpha = 32 ,    # Multiplier for LoRA 
         target_modules = [ "q_proj" ,   "k_proj" ,   "v_proj" ,   "o_proj" ],    # Target modules for optimization 
         lora_dropout = 0.05 ,    # Exclusion share for LoRA 
         bias = "none" ,    # Offset Setting 
         task_type = "CAUSAL_LM"    # Task type 
     ) 


     model   =   get_peft_model ( model ,   config )    # Applying PEFT optimization to a model 
     print ( f 'local rank =  { local_rank } , rank =  { rank } ' )    # Displaying information about the rank of a process 
     device   =   torch . device ( f 'cuda: { local_rank } ' )    # Setting the device for the current process 
     model . to ( device )    # Moving a model to a device 
     model   =   DDP ( model ,   device_ids = [ local_rank ],   find_unused_parameters = False )    # Model wrapper for distributed learning 


     # Loading and preparing data 
     # data = load_dataset("csv", data_files=BASE_DIR + "/midjourney_prompt_dataset.csv")  # Loading a dataset 
     data   =   load_dataset ( "bittu9988/mid_journey_prompts" ) 


     # Function for generating and tokenizing requests 
     def   generate_and_tokenize_prompt ( data_point ): 
         full_prompt   =   f """ 
         <human>:  { data_point [ "User" ] } 
         <assistant>:  { data_point [ "Prompt" ] } 
         """ . strip ()    # Generating a complete request 
         tokenized_full_prompt   =   tokenizer ( full_prompt ,   padding = True ,   truncation = True )    # Request Tokenization 
         return   tokenized_full_prompt 


     tokenized_data   =   data [ "train" ] . shuffle () . map ( generate_and_tokenize_prompt )    # Tokenization and data shuffling 
     tokenized_data   =   tokenized_data . remove_columns ( data [ "train" ] . column_names )    # Removing unnecessary columns from data 
     # Settings for training 
     training_args   =   TrainingArguments ( 
         per_device_train_batch_size = 1 ,    # Batch size per device 
         gradient_accumulation_steps = 8 ,    # Gradient accumulation steps 
         num_train_epochs = epochs ,    # Number of training epochs 
         learning_rate = 2e-4 ,    # Learning rate 
         fp16 = True ,    # Using 16-bit floating points 
         save_total_limit = 3 ,    # Limit on the number of model saves 
         logging_steps = 1 ,    # Logging steps 
         output_dir = "experiments" ,    # Output directory 
         optim = "paged_adamw_8bit" ,    # Optimizer 
         lr_scheduler_type = "cosine" ,    # Learning Rate Scheduler Type 
         warmup_ratio = 0.05 ,    # Heating proportion 
         remove_unused_columns = False ,    # Removing unused columns 
         local_rank = local_rank    # Local rank 
     ) 


     # Initializing the trainer to train the model 
     trainer   =   Trainer ( 
         model = model , 
         train_dataset = tokenized_data , 
         args = training_args , 
         data_collator = DataCollatorForLanguageModeling ( tokenizer ,   mlm = False ), 
     ) 


     trainer . train ()    # Start training 
     torch . save ( model . state_dict (),   BASE_DIR   +   f "/model.bin" )    # Saving the trained model 
     cleanup ()    # Resource Cleanup
Была ли статья полезной ?
Предыдущая статья
Общий алгоритм работы с PyTorch DDP
Следующая статья
Руководства по работе с Jupyter Server
/ DOCS
ML Space