add zh docs for tf32

.gitignore

@@ -144,4 +144,4 @@ core.*
 # Run
 llm_ckpts
 events.*
-memory_trace
+memory_trace*

doc/code-docs/source/index.rst

@@ -47,6 +47,14 @@ InternLM
 
    parallel
 
+TF32训练
+-------------------
+
+.. toctree::
+   :maxdepth: 2
+
+   tf32
+
 模型备份
 --------------------
 

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