LNO / training.py

Upload 22 files

3284d90 verified over 1 year ago

11.7 kB

import datasets
datasets.config.DOWNLOADED_DATASETS_PATH = "/mnt/jeff/huggingface/data"
import os
os.environ['HF_HOME'] = '/mnt/jeff/huggingface'

import argparse
import json
import os
from pathlib import Path

import numpy as np
import torch
import sacrebleu

from datasets import load_dataset
from torch.utils.data import Dataset, ConcatDataset
from tqdm import tqdm
from transformers import (
    AutoProcessor,
    AutoModel,
    BatchFeature,
    Trainer,
    TrainingArguments,
    StoppingCriteria,
    StoppingCriteriaList,
)
from collections import defaultdict

import soundfile as sf
from datasets import Audio
import random
from ASRDataset import *


def count_parameters_by_module(model):
    # dictionary for parameters number by modules
    module_params = defaultdict(lambda: {"total": 0, "trainable": 0})
    
    # all params
    total_params = 0
    total_trainable_params = 0
    
    # Check Embedding Token masks
    embedding_masks = {}
    for name, param in model.named_parameters():
        if 'embed_tokens.weight' in name and hasattr(param, '_backward_hooks') and param._backward_hooks:
            # check if params has embedding_grad_mask_hook
            for hook_id, hook_fn in param._backward_hooks.items():
                if hook_fn.__code__.co_name == 'embedding_grad_mask_hook':
                    # Accessing mask variables in the closure of hook functions
                    for cell in hook_fn.__closure__ or []:
                        if isinstance(cell.cell_contents, torch.Tensor) and cell.cell_contents.dtype == torch.bool:
                            # check mask tensor
                            embedding_masks[name] = ~cell.cell_contents  # True : Trainable
                 
    # Count params by modules
    for name, param in model.named_parameters():
        # extracts top module_name
        module_name = name.split('.')[0]
        param_count = param.numel()
        
        module_params[module_name]["total"] += param_count
        total_params += param_count
        
        if param.requires_grad:
            # Only count for real trainable params. (with masks)
            if name in embedding_masks:
                trainable_count = embedding_masks[name].sum().item()
                module_params[module_name]["trainable"] += trainable_count
                total_trainable_params += trainable_count
            else:
                module_params[module_name]["trainable"] += param_count
                total_trainable_params += param_count
    
    print(f"All Params: {total_params:,}")
    print(f"Trainable Params: {total_trainable_params:,} ({total_trainable_params/total_params*100:.2f}%)")
    print("\nParams by Module:")
    
    for module_name, counts in sorted(module_params.items()):
        trainable_percentage = counts["trainable"] / counts["total"] * 100 if counts["total"] > 0 else 0
        total_percentage = counts["total"] / total_params * 100
        
        print(f"- {module_name}:")
        print(f"  Total: {counts['total']:,} ({total_percentage:.2f}% of model)")
        print(f"  Trainable: {counts['trainable']:,} ({trainable_percentage:.2f}% of module)")
    
    return module_params

def create_model(model_name_or_path, revision="main", use_flash_attention = False):
    model = AutoModel.from_pretrained(
        model_name_or_path,
        revision=revision,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        attn_implementation="flash_attention_2" if use_flash_attention else "eager",
        trust_remote_code=True,
    )
    
    # Set use_cache to False after model loaded
    model.config.use_cache = False

    # Freeze all parameters
    for param in model.parameters():
        param.requires_grad = False
    
    model.set_lora_adapter('speech')
    model.to(torch.bfloat16)
    
    # (Optional) unfreeze audio_tower parameters
    for param in model.audio_tower.parameters():
       param.requires_grad = True

    # Only unfreeze audio_projector parameters
    for param in model.audio_projector.parameters():
        param.requires_grad = True

    # (Optional) unfreeze audio embed_tokens
    train_embed = True
    if train_embed:
        embed_tokens = model.language_model.model.model.embed_tokens
        
        embed_tokens.weight.requires_grad = False

        # Added Speech token IDs (only this tokens be trainable)
        trainable_token_ids = [128253, 128254]

        embed_tokens.weight.requires_grad = True
        mask = torch.ones_like(embed_tokens.weight, dtype=torch.bool)
        mask[trainable_token_ids] = False  # Trainable Tokens are False (unfreeze), else True (freeze)

        # backward hook, with gradient masking
        def embedding_grad_mask_hook(grad):
            return grad.masked_fill(mask, 0)

        embed_tokens.weight.register_hook(embedding_grad_mask_hook)

        model.language_model.model.model.embed_tokens = embed_tokens
        
    count_parameters_by_module(model)

    return model

model_name_or_path = '/mnt/jeff/InCar/LlamaNemotronOmni/Llama-NemotronOmni'
use_flash_attention = False

output_dir = '../nemotron_tmp3'
batch_size = 128
batch_size_per_gpu = 4
learning_rate = 4.0e-5 # 1.0e-4 for fine-tuning
wd = 0.01
num_train_epochs = 10

revision = "main" #"v1.0"

processor = AutoProcessor.from_pretrained(
    model_name_or_path,
    revision=revision,
    trust_remote_code=True,
)

model = create_model(
    model_name_or_path,
    revision=revision,
    use_flash_attention=use_flash_attention,
)

train_datasets = []

custom_yating_tw_loc = TWCostumData(processor=processor,
                      csv_path='/mnt/jeff/InCar/data/tw_data/taiwan_location-srdc_tts-20250505-yating-1-2s-breezyvoice.csv')
train_datasets.append(custom_yating_tw_loc)  # 9458

custom_tw_loc = TWCostumData(processor=processor,
                      csv_path='/mnt/jeff/InCar/data/tw_data/taiwan_location-srdc_tts-20250509-common_voice_16_1-TW.csv')
train_datasets.append(custom_tw_loc) # 1500

custom_tw_loc2 = TWCostumData(processor=processor,
                      csv_path='/mnt/jeff/InCar/data/tw_data/taiwan_location-srdc_tts-20250529-common_voice_16_1-TW.csv')
train_datasets.append(custom_tw_loc2) # 9458

# custom_yating_tw_road = TWCostumData(processor=processor,
#                       csv_path='/mnt/jeff/InCar/data/tw_data/taiwan_road-srdc_tts-20250430-yating-1-2s-breezyvoice.csv')
# train_datasets.append(custom_yating_tw_road) # 35224

custom_tw_road = TWCostumData(processor=processor,
                      csv_path='/mnt/jeff/InCar/data/tw_data/taiwan_road-srdc_tts-20250509-common_voice_16_1-TW.csv')
train_datasets.append(custom_tw_road) # 1500

custom_tw_road2 = TWCostumData(processor=processor,
                      csv_path='/mnt/jeff/InCar/data/tw_data/taiwan_road-srdc_tts-20250529-common_voice_16_1-TW.csv')
train_datasets.append(custom_tw_road2) # 35224



# common voice asr
commonvoice_speech_tw2 = CommonVoiceDataset(  # 45689
    processor=processor,
    source_lang="zh-TW",
    split="other"
)
train_datasets.append(commonvoice_speech_tw2)

commonvoice_speech_cn = CommonVoiceDataset(  #28868
    processor=processor,
    source_lang="zh-CN",
    split="train[:10%]"
)
train_datasets.append(commonvoice_speech_cn)


commonvoice_speech_tw = CommonVoiceDataset(  #6812
    processor=processor,
    source_lang="zh-TW",
    split="train"
)
train_datasets.append(commonvoice_speech_tw)




# Libri Speech Clean ASR mode (English -> English text)
libri_speech_clean = LibriSpeechDataset(  #103781
    processor=processor,
    subset="clean",
    split="train.360"
)
train_datasets.append(libri_speech_clean)


# Fleurs ASR mode (English -> English text)
en_asr_fleurs = FleursDataset(  #2550
    processor=processor,
    split="train",
    source_lang="en_us",  # English
    mode="asr"
)
train_datasets.append(en_asr_fleurs)

ch_asr_fleurs = FleursDataset(  #3152
    processor=processor,
    split="train",
    source_lang="cmn_hans_cn",
    mode="asr"
)
train_datasets.append(ch_asr_fleurs)





# en_ch_ast_fleurs = FleursDataset(
#     processor=processor,
#     split="train",
#     source_lang="en_us",
#     target_lang="cmn_hans_cn",
#     mode="ast"
# )
# train_datasets.append(en_ch_ast_fleurs)

# ch_en_ast_fleurs = FleursDataset(
#     processor=processor,
#     split="train",
#     source_lang="cmn_hans_cn",
#     target_lang="en_us",
#     mode="ast"
# )
# train_datasets.append(ch_en_ast_fleurs)

print("Count Num of Datasets", len(train_datasets))
print([len(dataset) for dataset in train_datasets])

# ConcatDataset
train_dataset = ConcatDataset(train_datasets) if len(train_datasets) > 1 else train_datasets[0]
print("Count Length of Datas", len(train_dataset))



# Check GPUs
num_gpus = torch.cuda.device_count()
print(f'training on {num_gpus} GPUs')

assert (
    batch_size % (num_gpus * batch_size_per_gpu) == 0
), 'Batch size must be divisible by the number of GPUs'
gradient_accumulation_steps = batch_size // (num_gpus * batch_size_per_gpu)

# hard coded training args
dp_config = {
    "fp16": {
      "enabled": "auto",
      "loss_scale": 0,
      "loss_scale_window": 1000,
      "initial_scale_power": 16,
      "hysteresis": 2,
      "min_loss_scale": 1
  },
   "zero_optimization": {
       "stage": 2,
       "allgather_partitions": True,
       "allgather_bucket_size": 5e8,
       "overlap_comm": False,
       "reduce_scatter": True,
       "reduce_bucket_size": 5e8,
       "contiguous_gradients": True,
       "cpu_offload": True
   },

   "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "optimizer": {
        "type": "AdamW",
        "params": {
          "lr": "auto",
          "betas": 'auto',
          "eps": 'auto',
          "weight_decay": "auto"
        }
    },
    "scheduler": {
      "type": "WarmupDecayLR",
      "params": {
          "warmup_min_lr": "auto",
          "warmup_max_lr": "auto",
          "warmup_num_steps": "auto",
          "total_num_steps": "auto"
      }
  },
    "gradient_clipping": 1.0,
    "zero_optimization": {
        "stage": 0
    }
}
training_args = TrainingArguments(
    num_train_epochs=num_train_epochs,
    per_device_train_batch_size=batch_size_per_gpu,
    gradient_checkpointing=False,
    gradient_checkpointing_kwargs={'use_reentrant': False},
    gradient_accumulation_steps=gradient_accumulation_steps,
    optim='adamw_torch',
    adam_beta1=0.9,
    adam_beta2=0.95,
    adam_epsilon=1e-7,
    learning_rate=learning_rate,
    weight_decay=wd,
    max_grad_norm=1.0,
    lr_scheduler_type='cosine',
    warmup_steps=50,
    logging_steps=10,
    output_dir=output_dir,
    save_total_limit=10,
    save_only_model=True,
    bf16=True,
    fp16=False,
    remove_unused_columns=False,
    report_to='none',
    deepspeed=None,
    disable_tqdm=False,
    dataloader_num_workers=16,
    save_strategy='epoch',
    # save_steps=2500,
    ddp_find_unused_parameters=True,

)

out_path = Path(training_args.output_dir)
out_path.mkdir(parents=True, exist_ok=True)

# create optimizer only for trainable params
optimizer = torch.optim.AdamW(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=learning_rate,
    weight_decay=wd,
    betas=(0.9, 0.95),
    eps=1e-7,
)

# Trainer Setting
trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=covost_collate_fn,
    train_dataset=train_dataset,
    optimizers=(optimizer, None)
)

trainer.train()


# # 1. Save LoRA Adapter
model.language_model.model.save_pretrained(output_dir)

# # 1-1. Delete Markdown file
# markdown_file = os.path.join(output_dir, "README.md")
# if os.path.exists(markdown_file):
#     os.remove(markdown_file)

# 2. Save entire model
model.save_pretrained(output_dir)

	import datasets
	datasets.config.DOWNLOADED_DATASETS_PATH = "/mnt/jeff/huggingface/data"
	import os
	os.environ['HF_HOME'] = '/mnt/jeff/huggingface'

	import argparse
	import json
	import os
	from pathlib import Path

	import numpy as np
	import torch
	import sacrebleu

	from datasets import load_dataset
	from torch.utils.data import Dataset, ConcatDataset
	from tqdm import tqdm
	from transformers import (
	AutoProcessor,
	AutoModel,
	BatchFeature,
	Trainer,
	TrainingArguments,
	StoppingCriteria,
	StoppingCriteriaList,
	)
	from collections import defaultdict

	import soundfile as sf
	from datasets import Audio
	import random
	from ASRDataset import *


	def count_parameters_by_module(model):
	# dictionary for parameters number by modules
	module_params = defaultdict(lambda: {"total": 0, "trainable": 0})

	# all params
	total_params = 0
	total_trainable_params = 0

	# Check Embedding Token masks
	embedding_masks = {}
	for name, param in model.named_parameters():
	if 'embed_tokens.weight' in name and hasattr(param, '_backward_hooks') and param._backward_hooks:
	# check if params has embedding_grad_mask_hook
	for hook_id, hook_fn in param._backward_hooks.items():
	if hook_fn.__code__.co_name == 'embedding_grad_mask_hook':
	# Accessing mask variables in the closure of hook functions
	for cell in hook_fn.__closure__ or []:
	if isinstance(cell.cell_contents, torch.Tensor) and cell.cell_contents.dtype == torch.bool:
	# check mask tensor
	embedding_masks[name] = ~cell.cell_contents # True : Trainable

	# Count params by modules
	for name, param in model.named_parameters():
	# extracts top module_name
	module_name = name.split('.')[0]
	param_count = param.numel()

	module_params[module_name]["total"] += param_count
	total_params += param_count

	if param.requires_grad:
	# Only count for real trainable params. (with masks)
	if name in embedding_masks:
	trainable_count = embedding_masks[name].sum().item()
	module_params[module_name]["trainable"] += trainable_count
	total_trainable_params += trainable_count
	else:
	module_params[module_name]["trainable"] += param_count
	total_trainable_params += param_count

	print(f"All Params: {total_params:,}")
	print(f"Trainable Params: {total_trainable_params:,} ({total_trainable_params/total_params*100:.2f}%)")
	print("\nParams by Module:")

	for module_name, counts in sorted(module_params.items()):
	trainable_percentage = counts["trainable"] / counts["total"] * 100 if counts["total"] > 0 else 0
	total_percentage = counts["total"] / total_params * 100

	print(f"- {module_name}:")
	print(f" Total: {counts['total']:,} ({total_percentage:.2f}% of model)")
	print(f" Trainable: {counts['trainable']:,} ({trainable_percentage:.2f}% of module)")

	return module_params

	def create_model(model_name_or_path, revision="main", use_flash_attention = False):
	model = AutoModel.from_pretrained(
	model_name_or_path,
	revision=revision,
	torch_dtype=torch.bfloat16,
	device_map="auto",
	attn_implementation="flash_attention_2" if use_flash_attention else "eager",
	trust_remote_code=True,
	)

	# Set use_cache to False after model loaded
	model.config.use_cache = False

	# Freeze all parameters
	for param in model.parameters():
	param.requires_grad = False

	model.set_lora_adapter('speech')
	model.to(torch.bfloat16)

	# (Optional) unfreeze audio_tower parameters
	for param in model.audio_tower.parameters():
	param.requires_grad = True

	# Only unfreeze audio_projector parameters
	for param in model.audio_projector.parameters():
	param.requires_grad = True

	# (Optional) unfreeze audio embed_tokens
	train_embed = True
	if train_embed:
	embed_tokens = model.language_model.model.model.embed_tokens

	embed_tokens.weight.requires_grad = False

	# Added Speech token IDs (only this tokens be trainable)
	trainable_token_ids = [128253, 128254]

	embed_tokens.weight.requires_grad = True
	mask = torch.ones_like(embed_tokens.weight, dtype=torch.bool)
	mask[trainable_token_ids] = False # Trainable Tokens are False (unfreeze), else True (freeze)

	# backward hook, with gradient masking
	def embedding_grad_mask_hook(grad):
	return grad.masked_fill(mask, 0)

	embed_tokens.weight.register_hook(embedding_grad_mask_hook)

	model.language_model.model.model.embed_tokens = embed_tokens

	count_parameters_by_module(model)

	return model

	model_name_or_path = '/mnt/jeff/InCar/LlamaNemotronOmni/Llama-NemotronOmni'
	use_flash_attention = False

	output_dir = '../nemotron_tmp3'
	batch_size = 128
	batch_size_per_gpu = 4
	learning_rate = 4.0e-5 # 1.0e-4 for fine-tuning
	wd = 0.01
	num_train_epochs = 10

	revision = "main" #"v1.0"

	processor = AutoProcessor.from_pretrained(
	model_name_or_path,
	revision=revision,
	trust_remote_code=True,
	)

	model = create_model(
	model_name_or_path,
	revision=revision,
	use_flash_attention=use_flash_attention,
	)

	train_datasets = []

	custom_yating_tw_loc = TWCostumData(processor=processor,
	csv_path='/mnt/jeff/InCar/data/tw_data/taiwan_location-srdc_tts-20250505-yating-1-2s-breezyvoice.csv')
	train_datasets.append(custom_yating_tw_loc) # 9458

	custom_tw_loc = TWCostumData(processor=processor,
	csv_path='/mnt/jeff/InCar/data/tw_data/taiwan_location-srdc_tts-20250509-common_voice_16_1-TW.csv')
	train_datasets.append(custom_tw_loc) # 1500

	custom_tw_loc2 = TWCostumData(processor=processor,
	csv_path='/mnt/jeff/InCar/data/tw_data/taiwan_location-srdc_tts-20250529-common_voice_16_1-TW.csv')
	train_datasets.append(custom_tw_loc2) # 9458

	# custom_yating_tw_road = TWCostumData(processor=processor,
	# csv_path='/mnt/jeff/InCar/data/tw_data/taiwan_road-srdc_tts-20250430-yating-1-2s-breezyvoice.csv')
	# train_datasets.append(custom_yating_tw_road) # 35224

	custom_tw_road = TWCostumData(processor=processor,
	csv_path='/mnt/jeff/InCar/data/tw_data/taiwan_road-srdc_tts-20250509-common_voice_16_1-TW.csv')
	train_datasets.append(custom_tw_road) # 1500

	custom_tw_road2 = TWCostumData(processor=processor,
	csv_path='/mnt/jeff/InCar/data/tw_data/taiwan_road-srdc_tts-20250529-common_voice_16_1-TW.csv')
	train_datasets.append(custom_tw_road2) # 35224



	# common voice asr
	commonvoice_speech_tw2 = CommonVoiceDataset( # 45689
	processor=processor,
	source_lang="zh-TW",
	split="other"
	)
	train_datasets.append(commonvoice_speech_tw2)

	commonvoice_speech_cn = CommonVoiceDataset( #28868
	processor=processor,
	source_lang="zh-CN",
	split="train[:10%]"
	)
	train_datasets.append(commonvoice_speech_cn)


	commonvoice_speech_tw = CommonVoiceDataset( #6812
	processor=processor,
	source_lang="zh-TW",
	split="train"
	)
	train_datasets.append(commonvoice_speech_tw)




	# Libri Speech Clean ASR mode (English -> English text)
	libri_speech_clean = LibriSpeechDataset( #103781
	processor=processor,
	subset="clean",
	split="train.360"
	)
	train_datasets.append(libri_speech_clean)


	# Fleurs ASR mode (English -> English text)
	en_asr_fleurs = FleursDataset( #2550
	processor=processor,
	split="train",
	source_lang="en_us", # English
	mode="asr"
	)
	train_datasets.append(en_asr_fleurs)

	ch_asr_fleurs = FleursDataset( #3152
	processor=processor,
	split="train",
	source_lang="cmn_hans_cn",
	mode="asr"
	)
	train_datasets.append(ch_asr_fleurs)





	# en_ch_ast_fleurs = FleursDataset(
	# processor=processor,
	# split="train",
	# source_lang="en_us",
	# target_lang="cmn_hans_cn",
	# mode="ast"
	# )
	# train_datasets.append(en_ch_ast_fleurs)

	# ch_en_ast_fleurs = FleursDataset(
	# processor=processor,
	# split="train",
	# source_lang="cmn_hans_cn",
	# target_lang="en_us",
	# mode="ast"
	# )
	# train_datasets.append(ch_en_ast_fleurs)

	print("Count Num of Datasets", len(train_datasets))
	print([len(dataset) for dataset in train_datasets])

	# ConcatDataset
	train_dataset = ConcatDataset(train_datasets) if len(train_datasets) > 1 else train_datasets[0]
	print("Count Length of Datas", len(train_dataset))



	# Check GPUs
	num_gpus = torch.cuda.device_count()
	print(f'training on {num_gpus} GPUs')

	assert (
	batch_size % (num_gpus * batch_size_per_gpu) == 0
	), 'Batch size must be divisible by the number of GPUs'
	gradient_accumulation_steps = batch_size // (num_gpus * batch_size_per_gpu)

	# hard coded training args
	dp_config = {
	"fp16": {
	"enabled": "auto",
	"loss_scale": 0,
	"loss_scale_window": 1000,
	"initial_scale_power": 16,
	"hysteresis": 2,
	"min_loss_scale": 1
	},
	"zero_optimization": {
	"stage": 2,
	"allgather_partitions": True,
	"allgather_bucket_size": 5e8,
	"overlap_comm": False,
	"reduce_scatter": True,
	"reduce_bucket_size": 5e8,
	"contiguous_gradients": True,
	"cpu_offload": True
	},

	"train_batch_size": "auto",
	"gradient_accumulation_steps": "auto",
	"optimizer": {
	"type": "AdamW",
	"params": {
	"lr": "auto",
	"betas": 'auto',
	"eps": 'auto',
	"weight_decay": "auto"
	}
	},
	"scheduler": {
	"type": "WarmupDecayLR",
	"params": {
	"warmup_min_lr": "auto",
	"warmup_max_lr": "auto",
	"warmup_num_steps": "auto",
	"total_num_steps": "auto"
	}
	},
	"gradient_clipping": 1.0,
	"zero_optimization": {
	"stage": 0
	}
	}
	training_args = TrainingArguments(
	num_train_epochs=num_train_epochs,
	per_device_train_batch_size=batch_size_per_gpu,
	gradient_checkpointing=False,
	gradient_checkpointing_kwargs={'use_reentrant': False},
	gradient_accumulation_steps=gradient_accumulation_steps,
	optim='adamw_torch',
	adam_beta1=0.9,
	adam_beta2=0.95,
	adam_epsilon=1e-7,
	learning_rate=learning_rate,
	weight_decay=wd,
	max_grad_norm=1.0,
	lr_scheduler_type='cosine',
	warmup_steps=50,
	logging_steps=10,
	output_dir=output_dir,
	save_total_limit=10,
	save_only_model=True,
	bf16=True,
	fp16=False,
	remove_unused_columns=False,
	report_to='none',
	deepspeed=None,
	disable_tqdm=False,
	dataloader_num_workers=16,
	save_strategy='epoch',
	# save_steps=2500,
	ddp_find_unused_parameters=True,

	)

	out_path = Path(training_args.output_dir)
	out_path.mkdir(parents=True, exist_ok=True)

	# create optimizer only for trainable params
	optimizer = torch.optim.AdamW(
	filter(lambda p: p.requires_grad, model.parameters()),
	lr=learning_rate,
	weight_decay=wd,
	betas=(0.9, 0.95),
	eps=1e-7,
	)

	# Trainer Setting
	trainer = Trainer(
	model=model,
	args=training_args,
	data_collator=covost_collate_fn,
	train_dataset=train_dataset,
	optimizers=(optimizer, None)
	)

	trainer.train()


	# # 1. Save LoRA Adapter
	model.language_model.model.save_pretrained(output_dir)

	# # 1-1. Delete Markdown file
	# markdown_file = os.path.join(output_dir, "README.md")
	# if os.path.exists(markdown_file):
	# os.remove(markdown_file)

	# 2. Save entire model
	model.save_pretrained(output_dir)