[WIP] support tps#9513
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces token-level tracking and logging during training, including computing token counts in the data collator, accumulating them in the trainer state, and reporting token throughput in the progress logs. It also adds a reduction option to MeanMetric. However, several critical issues were identified in the review: an undefined variable reduction in MeanMetric.compute, missing imports and unsupported CPU tensor operations for distributed reduction in patcher.py, an incomplete statement (logs['']), a potential AttributeError when accessing self.state.num_tokens, and an inefficient list-flattening operation in the data collator.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| self.state, self.count = tensor[0].item(), int(tensor[1].item()) | ||
| if self.count == 0: | ||
| value = self.nan_value | ||
| if reduction == 'sum': |
There was a problem hiding this comment.
The variable reduction is not defined in the compute method. It should be accessed via self.reduction as initialized in __init__.
| if reduction == 'sum': | |
| if self.reduction == 'sum': |
| num_tokens = getattr(state, 'num_tokens', None) | ||
| if num_tokens is not None: | ||
| num_tokens = float(num_tokens) | ||
| if dist.is_initialized(): | ||
| num_tokens = torch.tensor(num_tokens) | ||
| dist.all_reduce(num_tokens, op=dist.ReduceOp.SUM) | ||
| tps = num_tokens / elapsed | ||
| logs['num_input_tokens_seen'] = round(num_tokens, 6) | ||
| logs['train_speed(tokens/s)'] = round(tps, 6) |
There was a problem hiding this comment.
There are several issues here:\n1. torch and dist are not imported in this file, which will cause a NameError at runtime.\n2. torch.tensor(num_tokens) creates a CPU tensor by default. In distributed training with NCCL backend, all_reduce on CPU tensors is not supported and will fail.\n3. num_tokens becomes a torch.Tensor after all_reduce, and calling round() on it or dividing it might cause issues or keep it as a tensor. It should be converted back to a float using .item().
num_tokens = getattr(state, 'num_tokens', None)\n if num_tokens is not None:\n import torch\n import torch.distributed as dist\n from swift.utils import get_current_device\n num_tokens = float(num_tokens)\n if dist.is_initialized():\n device = get_current_device()\n num_tokens_tensor = torch.tensor(num_tokens, device=device)\n dist.all_reduce(num_tokens_tensor, op=dist.ReduceOp.SUM)\n num_tokens = num_tokens_tensor.item()\n tps = num_tokens / elapsed\n logs['num_input_tokens_seen'] = round(num_tokens, 6)\n logs['train_speed(tokens/s)'] = round(tps, 6)| n_steps = state.global_step - self.current_step | ||
| num_tokens = logs.pop('num_tokens', None) | ||
| if num_tokens is not None and n_steps > 0: | ||
| logs[''] |
There was a problem hiding this comment.
This line is incomplete and will raise a KeyError at runtime because '' is not a valid key in logs. Please complete the logging logic or remove this line.
| if num_tokens is not None: | ||
| self.state.num_tokens += num_tokens |
There was a problem hiding this comment.
self.state (which is a TrainerState from transformers) does not have a num_tokens attribute by default. This will raise an AttributeError on the first training step. It should be initialized safely.
if num_tokens is not None:\n self.state.num_tokens = getattr(self.state, 'num_tokens', 0) + num_tokens| if self.packing and isinstance(batch[0], list): | ||
| batch = sum(batch, start=[]) | ||
| num_samples = len(batch) | ||
| num_tokens = sum(sum([b['lengths'] for b in batch], start=[])) |
There was a problem hiding this comment.
Using sum(..., start=[]) to flatten lists is highly inefficient ('lengths' key, it will raise a KeyError. A more efficient and robust approach is to sum the lengths generator-style.
| num_tokens = sum(sum([b['lengths'] for b in batch], start=[])) | |
| num_tokens = sum(sum(b.get('lengths', [])) for b in batch) |
1 participant
No description provided.