Update SMP v2 notebooks to use latest PT2.3.1-TSM2.4.0 release.

Author: viclzhu

training/distributed_training/pytorch/model_parallel_v2/gpt-neox/smp-finetuning-gpt-neox-fsdp-tp.ipynb

@@ -882,8 +882,8 @@
     "            }\n",
     "        },\n",
     "    },\n",
-    "    py_version=\"py310\",\n",
-    "    framework_version=\"2.2.0\",\n",
+    "    py_version=\"py311\",\n",
+    "    framework_version=\"2.3.1\",\n",
     "    # image_uri=$IMAGE,  # Either provide `framework_version` or `image_uri`\n",
     "    output_path=s3_output_bucket,\n",
     "    max_run=86400,\n",

training/distributed_training/pytorch/model_parallel_v2/gpt-neox/smp-train-gpt-neox-fsdp-tp.ipynb

@@ -873,8 +873,8 @@
     "            }\n",
     "        },\n",
     "    },\n",
-    "    py_version=\"py310\",\n",
-    "    framework_version=\"2.2.0\",\n",
+    "    py_version=\"py311\",\n",
+    "    framework_version=\"2.3.1\",\n",
     "    # image_uri=$IMAGE,  # Either provide `framework_version` or `image_uri`\n",
     "    output_path=s3_output_bucket,\n",
     "    max_run=86400,\n",
@@ -955,8 +955,8 @@
     "                }\n",
     "            },\n",
     "        },\n",
-    "        py_version=\"py310\",\n",
-    "        framework_version=\"2.2.0\",\n",
+    "        py_version=\"py311\",\n",
+    "        framework_version=\"2.3.1\",\n",
     "        # image_uri=$IMAGE,  # Either provide `framework_version` or `image_uri`\n",
     "        output_path=s3_output_bucket,\n",
     "        max_run=86400,\n",

training/distributed_training/pytorch/model_parallel_v2/llama_v2/smp-finetuning-llama-fsdp-tp.ipynb

@@ -867,8 +867,8 @@
     "            }\n",
     "        },\n",
     "    },\n",
-    "    py_version=\"py310\",\n",
-    "    framework_version=\"2.2.0\",\n",
+    "    py_version=\"py311\",\n",
+    "    framework_version=\"2.3.1\",\n",
     "    # image_uri=$IMAGE,  # Either provide `framework_version` or `image_uri`\n",
     "    output_path=s3_output_bucket,\n",
     "    max_run=86400,\n",

training/distributed_training/pytorch/model_parallel_v2/llama_v2/smp-train-llama-fsdp-tp-fp8.ipynb

@@ -831,8 +831,8 @@
     "            }\n",
     "        },\n",
     "    },\n",
-    "    py_version=\"py310\",\n",
-    "    framework_version=\"2.2.0\",\n",
+    "    py_version=\"py311\",\n",
+    "    framework_version=\"2.3.1\",\n",
     "    # image_uri=$IMAGE,  # Either provide `framework_version` or `image_uri`\n",
     "    output_path=s3_output_bucket,\n",
     "    max_run=86400,\n",
@@ -913,8 +913,8 @@
     "                }\n",
     "            },\n",
     "        },\n",
-    "        py_version=\"py310\",\n",
-    "        framework_version=\"2.2.0\",\n",
+    "        py_version=\"py311\",\n",
+    "        framework_version=\"2.3.1\",\n",
     "        # image_uri=$IMAGE,  # Either provide `framework_version` or `image_uri`\n",
     "        output_path=s3_output_bucket,\n",
     "        max_run=86400,\n",

training/distributed_training/pytorch/model_parallel_v2/mixtral/smp-train-mixtral-fsdp-ep.ipynb

@@ -916,8 +916,8 @@
     "                }\n",
     "            },\n",
     "        },\n",
-    "        py_version=\"py310\",\n",
-    "        framework_version=\"2.2.0\",\n",
+    "        py_version=\"py311\",\n",
+    "        framework_version=\"2.3.1\",\n",
     "        # image_uri=$IMAGE,  # Either provide `framework_version` or `image_uri`\n",
     "        output_path=s3_output_bucket,\n",
     "        max_run=86400,\n",