use SingletonMeta for _TensorboardAdapter

Author: zhuzilin

docs/en/get_started/quick_start.md

@@ -515,7 +515,7 @@ async def generate(args, sample: Sample, sampling_params) -> Sample:
     # 7. Fill and return Sample object
     sample.response = full_response
     sample.tokens = ...
-    sample.loss_masks = loss_masks
+    sample.loss_mask = loss_masks
     return sample
 ```
 

docs/zh/get_started/quick_start.md

@@ -521,7 +521,7 @@ async def generate(args, sample: Sample, sampling_params) -> Sample:
     # 7. 填充并返回 Sample 对象
     sample.response = full_response
     sample.tokens = ...
-    sample.loss_masks = loss_masks
+    sample.loss_mask = loss_masks
     return sample
 ```
 

examples/retool/generate_with_retool.py

@@ -300,7 +300,7 @@ async def generate(args, sample: Sample, sampling_params) -> Sample:
     sample.tokens = prompt_tokens_ids + response_token_ids
     sample.response_length = len(response_token_ids)
     sample.response = response
-    sample.loss_masks = loss_masks
+    sample.loss_mask = loss_masks
 
     # Store payload information for wandb logging
     sample.payload_text = prompt + response

slime/utils/tensorboard_utils.py

@@ -1,13 +1,14 @@
 import datetime
 import os
+from slime.utils.misc import SingletonMeta
 
 try:
     from torch.utils.tensorboard import SummaryWriter
 except:
     SummaryWriter = None
 
 
-class _TensorboardAdapter:
+class _TensorboardAdapter(metaclass=SingletonMeta):
     _instance = None
     _writer = None
 
@@ -22,20 +23,16 @@ class _TensorboardAdapter:
     # tb.log({"Accuracy": 0.9}, step=1)
     """
 
-    def __new__(cls, args):
+    def __init__(self, args):
         assert args.use_tensorboard, f"{args.use_tensorboard=}"
         tb_project_name = args.tb_project_name
         tb_experiment_name = args.tb_experiment_name
-        if cls._instance is None:
-            cls._instance = super(_TensorboardAdapter, cls).__new__(cls)
-            # Initialize if parameters are provided during first creation
-            if tb_project_name is not None or os.environ.get("TENSORBOARD_DIR", None):
-                if tb_project_name is not None and tb_experiment_name is None:
-                    tb_experiment_name = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-                cls._instance._initialize(tb_project_name, tb_experiment_name)
-            else:
-                raise ValueError("tb_project_name and tb_experiment_name, or TENSORBOARD_DIR are required")
-        return cls._instance
+        if tb_project_name is not None or os.environ.get("TENSORBOARD_DIR", None):
+            if tb_project_name is not None and tb_experiment_name is None:
+                tb_experiment_name = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+            self._instance._initialize(tb_project_name, tb_experiment_name)
+        else:
+            raise ValueError("tb_project_name and tb_experiment_name, or TENSORBOARD_DIR are required")
 
     def _initialize(self, tb_project_name, tb_experiment_name):
         """Actual initialization logic"""

train.py

@@ -11,11 +11,6 @@ def train(args):
     pgs = create_placement_groups(args)
     wandb_run_id = init_wandb_primary(args)
 
-    if args.use_tensorboard:
-        from slime.utils.tensorboard_utils import _TensorboardAdapter
-
-        _TensorboardAdapter(args)
-
     # create the rollout manager, with sglang engines inside.
     # need to initialize rollout manager first to calculate num_rollout
     rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"], wandb_run_id=wandb_run_id)