训练

训练命令

torchrun --master_addr=localhost --master_port=32555 \
    --nnodes=1 --nproc-per-node=2 --node_rank=0 \
    -m tzrec.train_eval \
    --pipeline_config_path multi_tower_din_taobao_local.config \
    --train_input_path data/taobao_data_train/\*.parquet \
    --eval_input_path data/taobao_data_eval/\*.parquet \
    --model_dir experiments/multi_tower_din_taobao_local
  • –pipeline_config_path: 训练用的配置文件

  • –model_dir: 模型训练目录

  • –train_input_path: 训练数据的输入路径,可以支持MaxCompute表/Parquet/CSV等路径,详见文档

  • –eval_input_path: 评估数据的输入路径

  • –continue_train: 是否增量训练

  • –fine_tune_checkpoint: 增量训练的checkpoint路径,如experiments/multi_tower_din_taobao_local/model.ckpt-0,如果不设置,增量训练使用model_dir下最近的检查点

  • –edit_config_json: 命令行以json的方式动态修改配置文件,如{“model_dir”:“experiments/”,“feature_configs[0].raw_feature.boundaries”:[4,5,6,7]}

  • –ignore_restore_optimizer: 加载checkpoint参数时,忽略加载优化器的参数

环境变量

  • ODPS_ENDPOINT: 在PAI-DLC/PAI-DSW环境,数据为MaxCompute表的情况下需设置,详见文档的OdpsDataset章节

  • ODPS_CONFIG_FILE_PATH: 在本地环境,数据为MaxCompute表的情况下需设置为odps_conf的路径,详见文档的OdpsDataset章节

训练配置

训练配置是指配置文件中的train_config,详细参考配置参考手册

  • sparse_optimizer: 稀疏参数(Embedding)的优化器和LR策略设置

sparse_optimizer {
    adagrad_optimizer {
        lr: 0.001
    }
    constant_learning_rate {
    }
}
  • dense_optimizer: 稠密参数的优化器LR策略设置

dense_optimizer {
    adam_optimizer {
        lr: 0.001
    }
    exponential_decay_learning_rate {
        decay_size: 1
        decay_factor: 0.7
        by_epoch: true
    }
}

LR策略可以支持按epoch更新或者按step更新

  • num_steps: 训练的步数,不能跟num_epochs同时设置

  • num_epochs: 训练的epoch数

  • save_checkpoints_steps: 保存模型的步数间隔,保存模型后会做一次评估

  • save_checkpoints_epochs: 保存模型的Epoch数间隔,保存模型后会做一次评估,与save_checkpoints_steps不能同时设置

  • fine_tune_checkpoint: 增量训练的checkpoint路径,也可以设置checkpoint目录,将会使用目录下最新的checkpoint

  • fine_tune_ckpt_var_map: 需要restore的参数列表文件路径,文件的每一行是{variable_name in current model}\t{variable name in old model ckpt}

    • 需要设置fine_tune_ckpt_var_map的情形:

      • 现在的模型和原有模型参数名不完全匹配,但想迁移,如embedding的名字不一样

        • 原有模型checkpoint的参数列表名可以通过如下命令获取

        python -m tzrec.tools.list_distcp_param --checkpoint_path experiments/multi_tower_din_taobao_local/model.ckpt-0
        
      • 现在的模型需要让模型加载一些预训练的Embedding参数

        • 预训练的参数可以通过如下方式转成tzrec的检查点

        # convert.py
        import torch
        from torch import distributed as dist
        from torch.distributed.checkpoint import save
        dist.init_process_group(backend='gloo')
        state_local = torch.load('pretrain.pth')
        state_dist = {'embedding.weight': state_local['embedding.weight']}
        save(state_dist, checkpoint_id="pretrain/model.ckpt-0"))
        
        torchrun --master_addr=localhost --master_port=32555 --nnodes=1 --nproc-per-node=1 --node_rank=0 convert.py
        
  • log_step_count_steps: 打印log和summary的步数间隔(如果打印时间间隔小于1s,会跳过打印)

  • is_profiling: 是否做训练性能分析,设置为true,会在模型目录下记录trace文件

  • use_tensorboard: 是否使用tensorboard,默认为true

  • tensorboard_summaries: 设置需要在tensorboard中展示的数据, 只在use_tensorboard为true时生效。 可选值为[“loss”, “learning_rate”, “parameter”, “global_gradient_norm”, “gradient_norm”, “gradient”], 默认值是[“loss”, “learning_rate”]。 在训练数据量大, 训练时长较长时, 酌情开启"parameter"、 “gradient”、“gradient_norm”, 避免出现性能问题。

  • cudnn_allow_tf32: cudnn是否打开tf32精度训练,默认为true

  • cuda_matmul_allow_tf32: cuda matmul 及 tzrec中的triton op是否打开tf32精度训练,默认为false,设置为true可加速训练

  • mixed_precision: 混合精度训练,默认不开启,可以配置’BF16’或’FP16’

  • grad_scaler: 梯度动态缩放配置,使用FP16的情况下建议配置,参考GradScaler

grad_scaler {
    init_scale: 65536
    growth_factor: 2
    backoff_factor: 0.5
    growth_interval: 2000
}
  • gradient_accumulation_steps: 梯度累计,累计gradient_accumulation_steps步梯度更新参数,适用于batch_size小的情况,默认为0,不开启

  • grad_clipping: 稠密梯度裁剪,用于防止稠密参数梯度爆炸,默认不开启

    • clipping_type: 裁剪类型,可选值为"norm"(按梯度范数裁剪)、“value”(按梯度值裁剪)、“none”(不裁剪),默认为"none"

    • max_gradient: 梯度裁剪阈值,默认为1.0

    • norm_type: 范数类型,默认为2.0(L2范数),可以设置为inf(无穷范数)

    • enable_global_grad_clip: 是否启用全局梯度裁剪(分布式训练),默认为false

grad_clipping {
    clipping_type: "norm"
    max_gradient: 1.0
    norm_type: 2.0
    enable_global_grad_clip: false
}

Note: 稀疏参数(Embedding)的梯度裁剪通过sparse_optimizer中各优化器(如adagrad_optimizer)的gradient_clippingmax_gradient字段配置,例如:

sparse_optimizer {
    adagrad_optimizer {
        lr: 0.001
        gradient_clipping: true
        max_gradient: 1.0
    }
}

训练性能优化

TorchEasyRec是以模型混合并行的方式进行训练的,会根据机间和卡间的通信拓扑环境的设置寻优最好的Embedding分片和并行计算的方式,在显存约束下最小化计算和通信开销

默认的机间的通信环境为RDMA,卡间的通信环境为NVLINK,如任务所运行的环境没有RDMA和NVLINK,需对如下环境变量做一些调整

对于使用的是不带NVLINK的机型,如ecs.gn7i-c32g1.32xlarge(4 * NVIDIA A10)ecs.gn6i-c24g1.24xlarge(4 * NVIDIA T4)等,需调小卡间带宽的环境变量INTRA_NODE_BANDWIDTH

export INTRA_NODE_BANDWIDTH=$(awk 'BEGIN {printf("%f", 30 * 1024 * 1024 * 1024 / 1000)}')

对于使用的是不带RDMA的机型,如ecs.gn7i-c32g1.32xlarge(4 * NVIDIA A10)ecs.gn6i-c24g1.24xlarge(4 * NVIDIA T4)ecs.gn6e-c12g1.24xlarge(8 * V100)等,需调小机间带宽的环境变量CROSS_NODE_BANDWIDTH,具体可以参考各机型的基础网络带宽文档并减去读数据所需网络带宽设置

export CROSS_NODE_BANDWIDTH=$(awk 'BEGIN {printf("%f", 3 * 1024 * 1024 * 1024 / 1000)}')

如果遇到CUDA OOM(out of memory)的情况,可以调高STORAGE_RESERVE_PERCENT适当增加显存的预留比例,默认值为0.15。调高后TorchEasyRec会在新的显存限制下重新寻优最好的Embedding分片和并行计算的方式

export STORAGE_RESERVE_PERCENT=0.5