Skip to main content

python

python param parser​

选中即为true的argument

parser.add_argument("--no-lora-compute", action="store_true")

ZeroMQ socket​

使用了 ZeroMQ（zmq） 库和 asyncio，来设置一个基于 ZeroMQ 的异步通信管道。它创建了两个 ZeroMQ 的 socket，一个用于 PUSH 类型的消息发送，一个用于 PULL 类型的消息接收。这些 socket 的作用是将消息在不同组件之间进行路由和传递。以下是代码的详细解释：

        # 创建了一个 ZeroMQ 异步上下文，用于管理 socket。Context(2) 表示该上下文下的 I/O 线程数为 2。
        context = zmq.asyncio.Context(2)
        # http server push -> pull router push -> pull detokenization push -> pull http server
        self.send_to_router = context.socket(zmq.PUSH)
        self.send_to_router.connect(f"tcp://127.0.0.1:{router_port}")

        self.recv_from_detokenization = context.socket(zmq.PULL)
        self.recv_from_detokenization.bind(f"tcp://127.0.0.1:{httpserver_port}")

torch.cuda.synchronize()​

torch.cuda.synchronize() 是 PyTorch 中的一个函数，用于在 GPU 上执行同步操作。在使用 GPU 进行计算时，许多操作是异步执行的，这样可以提高效率。但有时我们需要确保所有 GPU 操作都已经完成，尤其是在依赖前一个操作结果或进行性能测量时，这时 torch.cuda.synchronize() 就非常有用。

  • 同步 GPU 操作：它会阻塞当前代码的执行，直到所有在 GPU 上排队的计算操作完成。这意味着在 torch.cuda.synchronize() 之后，所有 GPU 操作都已完成，不再有未处理的操作。
  • 确保数据一致性：在某些情况下，后续代码依赖 GPU 上的计算结果，而 GPU 操作是异步执行的，所以需要 torch.cuda.synchronize() 来保证在继续后续操作之前所有计算都已完成。
  • 性能测试：在进行性能测量或记录执行时间时，通常需要使用 torch.cuda.synchronize() 来确保测量的时刻是所有 GPU 操作完成后的真实状态。

torch.cuda.synchronize() 在这段代码中用于确保 GPU 上的操作按顺序完成，特别是在执行模型合并和批处理填充时，确保数据一致性。通过这种方式，可以避免异步计算带来的潜在数据冲突，并保证后续操作在预期状态下执行。

await, async​

使用 async 声明的函数是一个 协程函数，调用它时不会立即执行，而是返回一个协程对象。async 函数中可以使用 await 暂停执行，等待异步任务完成。

await asyncio.gather(*rets)

  • 并发地运行多个 init_batch 调用。
  • 等待所有这些调用完成。
  • 如果这些调用返回值，它会将结果作为列表返回（尽管代码没有进一步处理结果）。
    async def _init_batch(self, batch: Batch):
        # nsys _init_batch
        torch.cuda.nvtx.range_push("_init_batch")
        reqs = [r.to_rpc_obj() for r in batch.reqs]
        rets = [self.model_rpcs[tp_rank].init_batch(batch.batch_id, reqs) for tp_rank in range(self.world_size)]
        await asyncio.gather(*rets)
        # nsys _init_batch
        torch.cuda.nvtx.range_pop()
        return

tqdm​

tqdm 是一个用于在 Python 中显示进度条的库，非常适合在循环或任务中跟踪进度。它可以直观地显示任务的完成程度、运行时间、估计剩余时间等信息，特别适合处理耗时操作，如数据处理、训练模型、文件下载等。

from tqdm import tqdm
import time

for i in tqdm(range(100)):
    time.sleep(0.1)  # 模拟耗时操作

输出示例：

80%|████████    | 80/100 [00:08<00:02,  9.99it/s]
  • 80%：当前完成的百分比。
  • 80/100：已完成的迭代次数与总次数。
  • [00:08<00:02]：当前用时和预计剩余时间。
  • 9.99it/s：每秒处理的迭代数。