如何使用VSCode调试基于torchrun启动的Llama2推理脚本？

我能通过命令行成功运行以下Llama2推理脚本，但在VSCode调试模式下尝试多次配置均失败：

torchrun --nproc_per_node 1 example_text_completion.py \
--ckpt_dir models/7B-Chat \
--tokenizer_path tokenizer.model \
--max_seq_len 128 --max_batch_size 4

尝试的调试配置及错误 #

配置1：使用torchrun作为模块 #

{
    "name": "Python: run_llama2_inference",
    "type": "python",
    "request": "launch",
    "module": "torchrun",
    "args": [
        "--nproc_per_node=1",
        "example_chat_completion.py",
        "--ckpt_dir=models/7B-Chat/",
        "--tokenizer_path=tokenizer.model",
        "--max_seq_len=512",
        "--max_batch_size=4",
    ],
    "console": "integratedTerminal",
    "justMyCode": true,
    "env": {
        "PYTHONPATH": "${workspaceFolder}"
    }
},

错误信息：No module named torchrun

配置2：使用torch.distributed.launch作为模块 #

{
    "name": "Python: run_llama2_inference",
    "type": "python",
    "request": "launch",
    "module": "torch.distributed.launch",
    "args": [
        "--use-env",
        "example_chat_completion.py",
        "--nproc_per_node=1",
        "--ckpt_dir=models/7B-Chat/",
        "--tokenizer_path=tokenizer.model",
        "--max_seq_len=512",
        "--max_batch_size=4",
    ],
    "console": "integratedTerminal",
    "justMyCode": true,
    "env": {
        "PYTHONPATH": "${workspaceFolder}"
    }
},

错误信息：Could not consume arg: --nproc_per_node=1

解决方案 #

方案1：直接调试脚本，手动注入分布式环境变量 #

跳过torchrun封装，直接调试目标脚本，通过环境变量模拟单节点分布式环境：

{
    "name": "Python: Llama2 Inference Debug",
    "type": "python",
    "request": "launch",
    "program": "${workspaceFolder}/example_chat_completion.py",
    "args": [
        "--ckpt_dir=models/7B-Chat/",
        "--tokenizer_path=tokenizer.model",
        "--max_seq_len=512",
        "--max_batch_size=4"
    ],
    "console": "integratedTerminal",
    "justMyCode": false,
    "env": {
        "PYTHONPATH": "${workspaceFolder}",
        "WORLD_SIZE": "1",
        "RANK": "0",
        "LOCAL_RANK": "0",
        "MASTER_ADDR": "127.0.0.1",
        "MASTER_PORT": "29500"
    }
}

方案2：用命令模式完全复刻终端调用 #

直接调用torchrun命令，确保和终端环境一致：

{
    "name": "Bash: Llama2 Inference Debug",
    "type": "bash",
    "request": "launch",
    "command": "${config:python.pythonPath} -m torchrun --nproc_per_node 1 example_chat_completion.py --ckpt_dir models/7B-Chat --tokenizer_path tokenizer.model --max_seq_len 512 --max_batch_size 4",
    "cwd": "${workspaceFolder}",
    "console": "integratedTerminal"
}

方案3：修复torchrun模块调用问题 #

torchrun是torch.distributed.run的别名，直接使用原名更可靠，同时确保VSCode使用的Python环境和终端一致：

{
    "name": "Python: torchrun Debug",
    "type": "python",
    "request": "launch",
    "module": "torch.distributed.run",
    "args": [
        "--nproc_per_node=1",
        "--standalone",
        "example_chat_completion.py",
        "--ckpt_dir=models/7B-Chat/",
        "--tokenizer_path=tokenizer.model",
        "--max_seq_len=512",
        "--max_batch_size=4"
    ],
    "console": "integratedTerminal",
    "justMyCode": true,
    "env": {
        "PYTHONPATH": "${workspaceFolder}"
    }
}

内容的提问来源于stack exchange，提问作者Yneedtobeserious