Training on the multi-GPUs but stuck in loss.backward()

username: winchest
cooked: <p>Hi! I ran my code on a single GPU and it worked well. But when I tried to run it on the server that has 2 GPUs, it hang on the loss.backward(). I use torch.nn.DataParallel to train on multi-GPUs. The train code is as follows:</p>
<pre><code class="lang-auto">def train_batch(
        model,
        optimizer,
        baseline,
        epoch,
        batch_id,
        step,
        batch,
        tb_logger,
        opts
):
    x, bl_val = baseline.unwrap_batch(batch)
    x = move_to(x, opts.device)
    bl_val = move_to(bl_val, opts.device) if bl_val is not None else None
    # Evaluate model, get costs and log probabilities
    cost, log_likelihood = model(x)
    # Evaluate baseline, get baseline loss if any (only for critic)
    bl_val, bl_loss = baseline.eval(x, cost) if bl_val is None else (bl_val, 0)
    # Calculate loss
    reinforce_loss = ((cost - bl_val) * log_likelihood).mean()
    loss = reinforce_loss + bl_loss

    # Perform backward pass and optimization step
    optimizer.zero_grad()
    print('1')
    with torch.autograd.set_detect_anomaly(True):
        print('in')
        loss.backward()
        print('out')
    print('2')
    # Clip gradient norms and get (clipped) gradient norms for logging
    grad_norms = clip_grad_norms(optimizer.param_groups, opts.max_grad_norm)
    print('3')
    optimizer.step()
    print('7')
</code></pre>
<p>The output is:<br>
1<br>
in<br>
and maintain this situation.<br>
The state of the GPUs:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/f/dfe39c33f63ca039f234072d4715f141bcc5552e.png" data-download-href="https://discuss.pytorch.org/uploads/default/dfe39c33f63ca039f234072d4715f141bcc5552e" title="屏幕截图 2023-03-11 105500"><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/f/dfe39c33f63ca039f234072d4715f141bcc5552e.png" alt="屏幕截图 2023-03-11 105500" data-base62-sha1="vWC7XKe1RPrTS0Ya79qWdgVcQ3Q" width="690" height="353" data-dominant-color="131313"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">屏幕截图 2023-03-11 105500</span><span class="informations">886×454 13.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Pytorch: 1.12.1<br>
Python: 3.8.11<br>
GPU: Nvidia GeForce RTX 2080 Ti</p>

username: ptrblck
cooked: <p>Was this setup working before with multi-GPUs or is it new?<br>
In the latter case, could you check if disabling IOMMU might help as described <a href="https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/troubleshooting.html#pci-access-control-services-acs">here</a>?</p>

username: winchest
cooked: <p>Thanks for your reply!  But I’m a liitle confused about what u say <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=12" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20">. Do u mean that whether the train with the same setting worked well on multi-GPUs? Actually, I have an original code with the same setting  (GPU,Pytorch,Python) and it worked well. But then I made some change (didn’t change the train code) amd it got problems.</p>

username: ptrblck
cooked: <p>What kind of changes did you do that caused the issues?</p>

username: winchest
cooked: <p>The original work is a seq2seq model. I just add some states and change the mask rule. But it worked well on the single GPU. Could the added variables cause the issue?</p>

username: winchest
cooked: <p>I ran the command <code>sudo lspci -vvv | grep ACSCtl</code>and it seems that ASC is not enabled.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/5/b55804dc49f42700e094efac60b90bfd4e7c004f.png" data-download-href="https://discuss.pytorch.org/uploads/default/b55804dc49f42700e094efac60b90bfd4e7c004f" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/5/b55804dc49f42700e094efac60b90bfd4e7c004f.png" alt="image" data-base62-sha1="pSf07x4AeeCafikCqAQFygYuA3R" width="690" height="92" data-dominant-color="232222"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1145×153 3.69 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: ptrblck
cooked: <p>I don’t think IOMMU is at fault since the original code was already working while your changes seem to have broken the code as you have explained.<br>
Therefore, I would recommend checking each change separately to narrow down which exactly broke the code.</p>

username: winchest
cooked: <p>Thank you for your reply! I considered your suggestion and I found that two new inputs caused the hang. But I still son;t know why this happened. The code include the two inputs (ltw, rtw )is down below.</p>
<pre><code class="lang-auto">class StateCVRP(NamedTuple):
    # Fixed input
    coords: torch.Tensor
    demand: torch.Tensor
    ltw: torch.Tensor  
    rtw: torch.Tensor  
    ids: torch.Tensor  # Keeps track of original fixed data index of rows

    # State
    prev_a: torch.Tensor  
    used_capacity: torch.Tensor   
    visited_: torch.Tensor  # Keeps track of tasks that have been visited
    lengths: torch.Tensor  
    cur_index:torch.Tensor
    i: torch.Tensor  # Keeps track of step
    VEHICLE_CAPACITY = 1.0  # Hardcoded
    VEHICLE_V = 3   

    def initialize(input, visited_dtype=torch.uint8):

        demand = input['demand']
        depot = input['depot']
        loc_id = input['loc_id']
        ltw = input['ltw']
        rtw = input['rtw']

        batch_size, n_loc, _ = loc_id.size()  
        return StateCVRP(
            coords=torch.cat((depot[:, None, :], loc_id), -2).type(torch.long),
            demand=demand,
            ltw=ltw,
            rtw=rtw,
            ids=torch.arange(batch_size, dtype=torch.int64, device=loc_id.device)[:, None],  # Add steps dimension
            prev_a=torch.zeros([batch_size, 10], dtype=torch.long, device=loc_id.device),
            used_capacity=demand.new_zeros(batch_size, 10, 1),
            cur_index=input['depot'][:,None,:].type(torch.long),  # 1024，1,2
            visited_=(  
                torch.zeros(
                    batch_size, 1, n_loc + 1,
                    dtype=torch.uint8, device=loc_id.device
                )
                if visited_dtype == torch.uint8
                else torch.zeros(batch_size, 1, (n_loc + 63) // 64, dtype=torch.int64, device=loc_id.device)  
            ),
            lengths=torch.zeros(batch_size, 1, device=loc_id.device),
            i=torch.zeros(1, dtype=torch.int64, device=loc_id.device)  # Vector with length num_steps
        )
</code></pre>
<pre><code class="lang-auto">def make_instance(args):
    depot = args['depot']
    loc_id = args['loc_id']
    demand = args['demand']
    ltw = args['ltw']
    rtw = args['rtw']
    capacity = args['CAPACITIES']
    grid_size = 1

    return {
        'loc_id': torch.tensor(loc_id, dtype=torch.float) / grid_size,
        'demand': torch.tensor(demand, dtype=torch.float) / capacity,
        'depot': torch.tensor(depot, dtype=torch.float) / grid_size,
        'ltw' : torch.tensor(ltw, dtype=torch.float) / grid_size,
        'rtw' : torch.tensor(rtw, dtype=torch.float) / grid_size,
    }
</code></pre>
<pre><code class="lang-auto">class VRPDataset(Dataset):
    
    def __init__(self, filename=None, size=50, num_samples=1000000, offset=0, distribution=None):
        super(VRPDataset, self).__init__()

        self.data_set = []
        if filename is not None:
            assert os.path.splitext(filename)[1] == '.pkl'

            with open(filename, 'rb') as f:
                data = pickle.load(f)
                a = data[offset:offset+num_samples]
            self.data = [make_instance(args) for args in data[offset:offset+num_samples]]

        else:

            CAPACITIES = {
                10: 20.,
                20: 30.,
                50: 40.,
                100: 50.
            }

            self.data = [
                {
                    'depot': torch.full([2], 2, dtype=torch.float),
                    'loc_id': torch.randint(0, 57, (size,2)),
                    'demand': (torch.FloatTensor(size).uniform_(0, 9).int() + 1).float() / CAPACITIES[size],
                    'ltw': torch.randint(0, 100, (size,)),
                    'rtw': torch.randint(500, 600, (size,)),
                }
                for i in range(num_samples)
            ]

        self.size = len(self.data)
    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        return self.data[idx]
</code></pre>
<pre><code class="lang-auto"> def _init_embed(self, input):
     features = ('demand', 'ltw', 'rtw')
     return torch.cat(
         (
             self.init_embed_depot(input['depot'])[:, None, :],
             self.init_embed(torch.cat((
                 input['loc_id'],
                 *(input[feat][:, :, None] for feat in features)
             ), -1))
         ),
         1
     )
</code></pre>

username: winchest
cooked: <p>OMGGG！Finally I find that the os.environ code part seems to be conflict with torch.nn.DataParallel.So when I delete the os.environ code It worked!!</p>

username: ptrblck
cooked: <p>Could you share what exactly you have set via <code>os.environ</code> which apparently caused the hang?</p>

username: winchest
cooked: <p>It’s<br>
<code>os.environ['CUDA_LAUNCH_BLOCKING']='1'</code><br>
I added it before to debug but forget to delete it.</p>

username: yanhan-si
cooked: <p>I have stuck on this for hours! And this post saved me!</p>