[Fix] Merge previous contributions from fw/refactor to lite

arealite/README.md

@@ -737,4 +737,4 @@ dataloader = StatefulDataLoader(
 )
 for data in dataloader:
     assert isinstance(data, list)
-```
+```

arealite/api/cli_args.py

@@ -4,6 +4,9 @@
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 
+import uvloop
+
+uvloop.install()
 from hydra import compose as hydra_compose
 from hydra import initialize as hydra_init
 from omegaconf import MISSING, OmegaConf

pyproject.toml

@@ -53,9 +53,9 @@ dependencies = [
     "hydra-core==1.4.0.dev1",
     "packaging",
     "tabulate",
+    "gymnasium>=1.1.1",
     "torchdata",
     "autoflake",
-    "gymnasium",
     "tensordict",
 
     # Monitoring and logging

realhf/api/core/data_api.py

@@ -8,6 +8,7 @@
 import random
 import time
 from contextlib import contextmanager
+from functools import lru_cache
 
 # NOTE: We don't sue wildcard importing here because the type
 # `Sequence` has a very similar name to `SequenceSample`.
@@ -47,6 +48,7 @@
 RL_TASKS = ["math", "code", "rlhf", "stem"]
 
 
+@lru_cache(maxsize=8)
 def load_hf_tokenizer(
     model_name_or_path: str,
     fast_tokenizer=True,

requirements.txt

@@ -69,8 +69,8 @@ word2number
 Pebble
 timeout-decorator
 prettytable
+gymnasium>=1.1.1
 swanlab[dashboard]
 torchdata
 autoflake
-gymnasium
-tensordict
+tensordict