pytorch / tau Goto Github PK

---------- coverage: platform linux, python 3.7.12-final-0 -----------
Name                      Stmts   Miss  Cover
---------------------------------------------
pippy/IR.py                 371      7    98%
pippy/PipelineDriver.py     394    330    [16](https://github.com/jamesr66a/PiPPy/runs/5221199398?check_suite_focus=true#step:6:16)%
pippy/__init__.py             3      0   100%
pippy/version.py              2      2     0%
---------------------------------------------
TOTAL                       770    339    56%

$ test/launch_local_test_forward_backward.sh 
WARNING:torch.distributed.run:
*****************************************
Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
*****************************************
REPLICATE config: False -> MultiUseParameterConfig.TRANSMIT
GraphModule(
  (submod_0): GraphModule()
  (submod_1): GraphModule()
  (submod_2): GraphModule()
  (_loss): MSELoss()
)



def forward(self, x, target):
    submod_0 = self.submod_0(x)
    getitem_2 = submod_0[2]
    getitem = submod_0[0]
    getitem_1 = submod_0[1]
    submod_1 = self.submod_1(getitem, getitem_2)
    getitem_4 = submod_1[1]
    getitem_3 = submod_1[0]
    submod_2 = self.submod_2(getitem_3, getitem_1, getitem_4)
    _loss = self._loss(submod_2, target)
    stage_backward = pippy_IR_stage_backward(stage_output = _loss, output_grads = None, input_values = [submod_2, target]);  target = None
    getitem_5 = stage_backward[0]
    getitem_6 = stage_backward[1];  stage_backward = None
    getitem_7 = getitem_5[0]
    getitem_8 = getitem_5[1];  getitem_5 = None
    stage_backward_1 = pippy_IR_stage_backward(stage_output = submod_2, output_grads = getitem_7, input_values = [getitem_3, getitem_1, getitem_4]);  submod_2 = getitem_7 = getitem_3 = getitem_1 = getitem_4 = None
    getitem_9 = stage_backward_1[0]
    getitem_10 = stage_backward_1[1];  stage_backward_1 = None
    getitem_11 = getitem_9[0]
    getitem_12 = getitem_9[1]
    getitem_13 = getitem_9[2];  getitem_9 = None
    stage_backward_2 = pippy_IR_stage_backward(stage_output = submod_1, output_grads = [getitem_11, getitem_13], input_values = [getitem, getitem_2]);  submod_1 = getitem_11 = getitem_13 = getitem = getitem_2 = None
    getitem_14 = stage_backward_2[0]
    getitem_15 = stage_backward_2[1];  stage_backward_2 = None
    getitem_16 = getitem_14[0]
    getitem_17 = getitem_14[1];  getitem_14 = None
    stage_backward_3 = pippy_IR_stage_backward(stage_output = submod_0, output_grads = [getitem_16, getitem_12, getitem_17], input_values = [x]);  submod_0 = getitem_16 = getitem_12 = getitem_17 = x = None
    getitem_18 = stage_backward_3[0]
    getitem_19 = stage_backward_3[1];  stage_backward_3 = None
    getitem_20 = getitem_18[0];  getitem_18 = None
    sync_barrier = pippy_IR_sync_barrier(_loss, [getitem_6, getitem_10, getitem_15, getitem_19]);  _loss = getitem_6 = getitem_10 = getitem_15 = getitem_19 = None
    return sync_barrier
    
/fsx/users/jamesreed/pipeline_for_real/pippy/PipelineDriver.py:498: UserWarning: Running pipeline with 3 stages on world_size of 10. Remaining ranks will be idle.
  warnings.warn(f'Running pipeline with {len(executor_descriptors)} stages on world_size of {self.world_size}. '
Traceback (most recent call last):
  File "/fsx/users/jamesreed/pipeline_for_real/test/local_test_forward_backward.py", line 117, in <module>
    out = pipe_driver.run((input, target), {}, chunks=CHUNKS, _debug_mask_minibatches = DEBUG_MASK_MINIBATCHES)
  File "/fsx/users/jamesreed/pipeline_for_real/pippy/PipelineDriver.py", line 691, in run
    last_nodes.append(interp.run_until(lambda n: n.op == 'output'))
  File "/fsx/users/jamesreed/pipeline_for_real/pippy/PipelineDriver.py", line 606, in run_until
    self.env[node] = self.run_node(node)
  File "/fsx/users/jamesreed/pytorch/torch/fx/interpreter.py", line 152, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
  File "/fsx/users/jamesreed/pipeline_for_real/pippy/PipelineDriver.py", line 573, in call_function
    return args[0].remote().__getitem__(args[1])
  File "/fsx/users/jamesreed/pytorch/torch/distributed/rpc/rref_proxy.py", line 41, in _invoke_rpc
    rref_fut.wait()
RuntimeError: RPCErr:1:RPC ran for more than set timeout (60000 ms) and will now be marked with an error

import inspect

import transformers.utils.fx as fx
from pippy.IR import MultiUseParameterConfig, Pipe
from transformers import *

model = T5Model(T5Config())
print(model)

input_names = model.dummy_inputs.keys()
sig = inspect.signature(model.forward)
concrete_args = {p.name: p.default for p in sig.parameters.values() if p.name not in input_names}

hf_tracer = fx.HFTracer()

model_pipe = Pipe.from_tracing(model, MultiUseParameterConfig.TRANSMIT, tracer=hf_tracer,
                               concrete_args=concrete_args)

T5Model(
  (shared): Embedding(32128, 512)
  (encoder): T5Stack(...)
  (decoder): T5Stack(...)
  ...

Traceback (most recent call last):
  File "/Users/pbelevich/PycharmProjects/PiPPy/test/hf_t5_test.py", line 15, in <module>
    model_pipe = Pipe.from_tracing(model, MultiUseParameterConfig.TRANSMIT, tracer=hf_tracer,
  File "/Users/pbelevich/PycharmProjects/PiPPy/pippy/IR.py", line 606, in from_tracing
    return Pipe._from_traced(mod, traced, multi_use_param_spec, loss_fn, **kwargs)
  File "/Users/pbelevich/PycharmProjects/PiPPy/pippy/IR.py", line 457, in _from_traced
    mod_itr = getattr(mod_itr, atom)
  File "/Users/pbelevich/miniconda3/envs/PiPPy/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1186, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'GraphModule' object has no attribute 'decoder'

/home/pbelevich/local/PiPPy/pippy/PipelineDriver.py:393: UserWarning: Running pipeline with 13 stages on world_size of 20. Remaining ranks will be idle.
  warnings.warn(f'Running pipeline with {len(executor_descriptors)} stages on world_size of {self.world_size}. '
Traceback (most recent call last):
  File "/home/pbelevich/local/PiPPy/test/local_test_forward.py", line 55, in <module>
    out = pipe_driver.run(gpt2_input, chunks=5, _debug_mask_minibatches=True)
  File "/home/pbelevich/local/PiPPy/pippy/PipelineDriver.py", line 570, in run
    return self._retrieve_output_values(microbatch_interpreters, last_nodes, _debug_mask_minibatches, splits_per_arg)
  File "/home/pbelevich/local/PiPPy/pippy/PipelineDriver.py", line 605, in _retrieve_output_values
    sliced_outputs.append(result[start:end])
TypeError: unhashable type: 'slice'

        def forward(self, x):
            x = self.linear(x)
            pipe_split()
            pipe_split()
            pipe_split()
            return x

# Copyright (c) Meta Platforms, Inc. and affiliates
import argparse
import os
import socket

import torch
import torch.distributed.rpc as rpc
import torch.multiprocessing as mp

from pippy.IR import Pipe, pipe_split, TrivialLossWrapper, _null_coalesce_accumulate
from pippy.PipelineDriver import PipelineDriverFillDrain, PipelineDriver1F1B, PipelineDriverBase
from pippy.microbatch import TensorChunkSpec, CustomReducer

PROFILING_ENABLED = True
CHECK_NUMERIC_EQUIVALENCE = True

schedules = {
    'FillDrain': PipelineDriverFillDrain,
    '1F1B': PipelineDriver1F1B,
}

torch.fx.Tracer.proxy_buffer_attributes = True


def run_master(args):
    all_ranks = list(range(1, args.world_size))  # exclude master rank = 0
    chunks = len(all_ranks)
    batches = 1
    bs = 4 * chunks
    hid_dim = 50

    class Code(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.linear = torch.nn.Linear(hid_dim, hid_dim)

        def forward(self, x):
            x = self.linear(x)
            pipe_split()
            y = torch.relu(x)
            pipe_split()
            z = torch.sigmoid(x)
            pipe_split()
            return y + z

    c = Code()
    c.train()
    mse_loss = torch.nn.MSELoss()
    wrapper = TrivialLossWrapper(c, mse_loss)
    accum_pipe = Pipe.from_tracing(wrapper)
    assert 4 == len(list(accum_pipe.split_gm.children()))
    assert any(n.target == _null_coalesce_accumulate for n in accum_pipe.split_gm.graph.nodes)
    input = torch.randn(bs, hid_dim)
    target = torch.randn(bs, hid_dim)
    accum_pipe(input, target)

    args_chunk_spec = (TensorChunkSpec(0), TensorChunkSpec(0))
    kwargs_chunk_spec = {}
    output_chunk_spec = CustomReducer(torch.tensor(0.0), lambda a, b: a + b)
    pipe_driver: PipelineDriverBase = schedules[args.schedule](accum_pipe, args_chunk_spec, kwargs_chunk_spec,
                                                               output_chunk_spec, args.world_size - 1,
                                                               all_ranks=all_ranks, _debug_mask_minibatches=True)

    for i in range(batches):
        pipe_driver.run(chunks, input, target)


def run_worker(rank, world_size, args):
    print(f"rank = {rank} host/pid = {socket.gethostname()}/{os.getpid()}")
    os.environ['MASTER_ADDR'] = args.master_addr
    os.environ['MASTER_PORT'] = args.master_port
    if args.rank == -1:  # run via mp.spawn
        # each worker will see its GPU as `cuda:0`
        try:
            import subprocess
            device_count = int(subprocess.getoutput('nvidia-smi --list-gpus | wc -l'))
            os.environ['CUDA_VISIBLE_DEVICES'] = str(rank % device_count)
        except ValueError:
            pass
    assert not torch.cuda.is_available() or torch.cuda.device_count() == 1, \
        "Do not use torch.cuda.* before setting CUDA_VISIBLE_DEVICES"
    options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=256, _transports=["uv"])  # uv for AWS EFA instances
    if args.use_cuda and torch.cuda.is_available():
        for i in range(world_size):
            options.set_device_map(f"worker{i}", {0: 0})
    rpc.init_rpc(
        f"worker{rank}",
        rank=rank,
        world_size=world_size,
        rpc_backend_options=options
    )
    if rank == 0:
        run_master(args)
    rpc.shutdown()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--world_size', type=int, default=int(os.getenv("WORLD_SIZE", 5)))
    parser.add_argument('--rank', type=int, default=int(os.getenv("RANK", -1)))
    parser.add_argument('--master_addr', type=str, default=os.getenv('MASTER_ADDR', 'localhost'))
    parser.add_argument('--master_port', type=str, default=os.getenv('MASTER_PORT', '29500'))
    parser.add_argument('-s', '--schedule', type=str, default=list(schedules.keys())[0], choices=schedules.keys())
    parser.add_argument('--replicate', type=int, default=int(os.getenv("REPLICATE", '0')))
    parser.add_argument('--use_cuda', type=int, default=1)
    args = parser.parse_args()

    if args.rank == -1:
        mp.spawn(run_worker, args=(args.world_size, args,), nprocs=args.world_size, join=True)
    elif args.rank < args.world_size:
        run_worker(args.rank, args.world_size, args)
    else:
        print("I'm unused, exiting")

Traceback (most recent call last):
  File "/Users/pbelevich/miniconda3/envs/PiPPy/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/Users/pbelevich/PycharmProjects/PiPPy/examples/slurm/hf/t5/test_null_coalesce_accumulate.py", line 93, in run_worker
    run_master(args)
  File "/Users/pbelevich/PycharmProjects/PiPPy/examples/slurm/hf/t5/test_null_coalesce_accumulate.py", line 65, in run_master
    pipe_driver.run(chunks, input, target)
  File "/Users/pbelevich/PycharmProjects/PiPPy/pippy/PipelineDriver.py", line 871, in run
    last_nodes.append(interp.run_until(lambda n: n.op == 'output'))
  File "/Users/pbelevich/PycharmProjects/PiPPy/pippy/PipelineDriver.py", line 811, in run_until
    self.env[node] = super().run_node(node)
  File "/Users/pbelevich/miniconda3/envs/PiPPy/lib/python3.9/site-packages/torch/fx/interpreter.py", line 152, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
  File "/Users/pbelevich/PycharmProjects/PiPPy/pippy/PipelineDriver.py", line 800, in call_function
    raise AssertionError(f'Unknown operator {torch.typename(target)}')
AssertionError: Unknown operator pippy.IR._null_coalesce_accumulate

def forward(self, input_ids):
    ...
    _tensor_constant0 = self._tensor_constant0
    ...
    _tensor_constant0_1 = self._tensor_constant0
    ...
    _tensor_constant0_2 = self._tensor_constant0
    ...
    _tensor_constant0_3 = self._tensor_constant0
    ...
    _tensor_constant0_4 = self._tensor_constant0
    ...
    _tensor_constant0_5 = self._tensor_constant0
    ...
    _tensor_constant0_6 = self._tensor_constant0
    ...
    _tensor_constant0_7 = self._tensor_constant0
    ...
    _tensor_constant0_8 = self._tensor_constant0
    ...
    _tensor_constant0_9 = self._tensor_constant0
    ...
    _tensor_constant0_10 = self._tensor_constant0
    ...
    _tensor_constant0_11 = self._tensor_constant0
    ...

AttributeError: 'GraphModule' object has no attribute '_tensor_constant0'