diff --git a/.gitignore b/.gitignore index 8992a0f..236d8c9 100644 --- a/.gitignore +++ b/.gitignore @@ -144,4 +144,4 @@ core.* # Run llm_ckpts events.* -memory_trace +memory_trace* diff --git a/doc/code-docs/source/index.rst b/doc/code-docs/source/index.rst index c01ac54..f6b17c2 100644 --- a/doc/code-docs/source/index.rst +++ b/doc/code-docs/source/index.rst @@ -47,6 +47,14 @@ InternLM parallel +TF32训练 +------------------- + +.. toctree:: + :maxdepth: 2 + + tf32 + 模型备份 -------------------- diff --git a/doc/code-docs/source/tf32.rst b/doc/code-docs/source/tf32.rst new file mode 100644 index 0000000..06ad730 --- /dev/null +++ b/doc/code-docs/source/tf32.rst @@ -0,0 +1,45 @@ +TF32训练 +================== + +InternLM支持使用TF32训练模型。TensorFloat-32(TF32)是Nvidia在Ampere架构GPU上推出的专门运用于TensorCore的一种计算格式。其与其他常用数据格式的比较如下图: + +.. figure:: ../../imgs/tf32.png + :scale: 50% + :class: with-border + +使用TF32的前置条件: + +1. 输入数据类型为FP32,且计算为矩阵乘法及卷积相关运算,才可以使用TF32作为TensorCore的中间计算类型。 + + +2. Ampere架构的GPU。 + + +值得注意的是,TF32仅仅是在使用TensorCore时的一种中间计算格式,并不是一个完全的数据类型。因此,为了区分不同的精度与计算格式( ``BF16`` 、``FP16`` 、``FP32`` 、``TF32`` ),InternLM支持用户在 ``model config`` 中传入 ``torch.tf32`` 来表示想要使用TF32加速运算,本质上数据类型依旧为 ``FP32``。 + +.. code-block:: python + + model = dict( + checkpoint=False, # The proportion of layers for activation aheckpointing, the optional value are True/False/[0-1] + num_attention_heads=NUM_ATTENTION_HEAD, + embed_split_hidden=True, + vocab_size=VOCAB_SIZE, + embed_grad_scale=1, + parallel_output=True, + hidden_size=HIDDEN_SIZE, + num_layers=NUM_LAYER, + mlp_ratio=MLP_RATIO, + apply_post_layer_norm=False, + dtype="torch.tf32", # Support: "torch.float16", "torch.half", "torch.bfloat16", "torch.float32", "torch.tf32" + norm_type="rmsnorm", + layer_norm_epsilon=1e-5, + use_flash_attn=True, + num_chunks=1, # if num_chunks > 1, interleaved pipeline scheduler is used. + ) + +InternLM会根据 ``model config`` 中的 ``dtype`` 字符串来判断真正的数据类型。InternLM通过设置以下变量来开启TF32训练。 + +.. code-block:: python + + torch.backends.cudnn.allow_tf32 = True + torch.backends.cuda.matmul.allow_tf32 = True diff --git a/doc/imgs/tf32.png b/doc/imgs/tf32.png new file mode 100644 index 0000000..b7f8393 Binary files /dev/null and b/doc/imgs/tf32.png differ