[Fix] Fix TypedDict import error for Python < 3.12

Author: nil0x9

xtuner/v1/module/attention/attn_outputs.py

@@ -1,4 +1,4 @@
-from typing import TypedDict
+from typing_extensions import TypedDict
 
 import torch
 

xtuner/v1/ops/attn_imp.py

@@ -1,6 +1,6 @@
 import traceback
 from functools import lru_cache
-from typing import TypedDict
+from typing_extensions import TypedDict
 
 import torch
 import torch.nn as nn

xtuner/v1/rl/base/controller.py

@@ -1,6 +1,7 @@
 import math
 import os
-from typing import Literal, TypedDict
+from typing import Literal
+from typing_extensions import TypedDict
 
 import ray
 import torch

xtuner/v1/rl/base/worker.py

@@ -4,7 +4,7 @@
 import time
 from itertools import chain
 from pathlib import Path
-from typing import Dict, Iterable, List, TypeAlias, TypedDict, cast
+from typing import Dict, Iterable, List, TypeAlias, cast
 
 import ray
 import requests
@@ -16,7 +16,7 @@
 from ray.actor import ActorClass, ActorProxy
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 from torch.distributed.tensor import DTensor
-from typing_extensions import NotRequired
+from typing_extensions import NotRequired, TypedDict
 
 from transformers import AutoTokenizer
 from xtuner.v1.config.fsdp import FSDPConfig