| 107 | return ceil(num / mult) * mult |
| 108 | |
| 109 | def get_code_utilization(codes, codebook_size, get_global=False): |
| 110 | if get_global and dist.is_initialized(): |
| 111 | world_size = dist.get_world_size() |
| 112 | else: |
| 113 | world_size = 1 |
| 114 | |
| 115 | if world_size > 1: |
| 116 | gathered_tokens = [T.zeros_like(codes) for _ in range(world_size)] |
| 117 | dist.all_gather(gathered_tokens, codes) |
| 118 | gathered_tokens = T.cat(gathered_tokens, dim=0) |
| 119 | else: |
| 120 | gathered_tokens = codes |
| 121 | unique_tokens = len(T.unique(gathered_tokens)) |
| 122 | code_utilization = unique_tokens / min(gathered_tokens.numel(), codebook_size) |
| 123 | return code_utilization |
| 124 | |
| 125 | # tensor helpers |
| 126 | |