def register_common_model_topology_args()

in src/optimum/nvidia/utils/cli.py [0:0]

• 36 lines of code
• 1 McCabe index (conditional complexity)

def register_common_model_topology_args(parser: ArgumentParser) -> ArgumentParser:
    parser.add_argument(
        "--tensor-parallelism",
        type=int,
        default=1,
        dest="tp",
        help="Define the number of slice for each tensor each GPU will receive.",
    )
    parser.add_argument(
        "--pipeline-parallelism",
        type=int,
        default=1,
        dest="pp",
        help="Define the number of sections to split neural network layers",
    )
    parser.add_argument(
        "--world-size",
        type=int,
        default=1,
        help="Total number of GPUs over all the nodes.",
    )
    parser.add_argument(
        "--gpus-per-node",
        type=int,
        default=1,
        help="Total number of GPUs on a single node.",
    )
    parser.add_argument(
        "-o",
        "--opt-level",
        type=int,
        default=2,
        dest="optimization_level",
        help="Optimization level between 0 (no optimization) and 5 (maximum level of optimization).",
    )
    return parser