@MinasTyuru

I have tried in colab with TF version 2.2 and i am seeing the error message(ModuleNotFoundError: No module named 'tflight').Request you to share colab link or simple standaalone code with supporting files to reproduce the issue in our environment.It helps us in localizing the issue faster.Thanks!
`

@ravikyram Oops, sorry about that. I removed the line that was the problem and it should work now. Please keep in mind that the code must run on multiple workers to reproduce the issue. As far as I know, this is not possible on Google Colab.

Hi @MinasTyuru, thanks for providing very detailed debugging information. BTW you can use virtual GPUs in colab to simulate a multiworker environment, but NCCL is not supported in this case, so probably not helpful for debugging this particular issue.

Can you confirm that you do not see this issue if the CollectiveCommunication is set to AUTO? Additionally, can you try and run on TF nightly and let me know what happens?

@nikitamaia Thank you for the information! Yes, that's correct. Using AUTO or RING seems to work fine. Running on TF nightly may be a bit tricky on our infrastructure, I will get back to you on that.

@nikitamaia Hi Nikita, I was able to reproduce on TF-nightly 2.4.0.dev20200708. When using AUTO it works fine, but when using NCCL it immediately segfaults. Attached are the failed log with NCCL and the successful log with AUTO.

@MinasTyuru thank you for the report. It seems to be the case from the title, but I would like to confirm: does this issue only show up when using SyncBatchNormalization layer? (And goes away when not using it, or using a regular batch norm).

Getting a segfault is terrible, we will look into it.

We have used SyncBatchNormalization for some models before with multi worker, so I am trying to see what's different here since your repro is quite simple. Would you mind trying one thing: Can you make the batch size bigger? Currently seems 1. Usually we try to split the batch across the replicas but since batch size of 1 cannot be split, the other replicas will end up getting empty batches. I am wondering if that's what is triggering this. Although, it would not explain why it works for some steps and then fails..

After some more investigation, it looks like we can reproduce the issue in some contexts but not others, so it's possible that it is some sort of configuration issue. I thought that I was able to reproduce on TF-nightly, but I believe that was due to a mistake I made when installing CUDA, and now I think it may work fine on TF-nightly. It may be related to how we build TensorFlow from source or something of that nature. I will investigate some more and let you know if it seems to be a problem on our end or not.

thanks for the update @MinasTyuru .
@dubey is looking into the segfault more - he thinks there is a possible race condition which might be leading to the seg fault and he is looking into a possible fix.

Re: nightly, yes, your logs indicated that the issue with the nightly was that it was not able to find the GPUs and hence failed right away. I would guess that once you fix that issue, it may probably run into the segfault at some point later like with 2.2.

Yes, we are seeing something similar. When I used a debugger, it changed the behavior. For example, it seems like adding a breakpoint after line 176 of collective_nccl_reducer.cc, which introduces a very small delay, causes TensorFlow to immediately segfault. In this case, it seems that the col_ctx_ variable is set to a null pointer by bfc_allocator.cc. Perhaps this memory is inappropriately freed before it is finished being used? Is this behavior consistent with your hypothesis of the race condition? I'll post some more logs in ~12 hours.

We tried to reproduce on TF2.2 and TF-nightly, and on a clean build (i.e. installing TensorFlow from pip, etc.), it seems like it trains for a long time without any problems. But in our production environment (which changes some of the dependencies in the TensorFlow build and other build differences), we encounter the segfault, so perhaps subtle timing differences are triggering the underlying issue.

Attached are a stacktrace from when col_ctx_ is set to 0 by the memory allocator, and the assembly around that line. Thanks again!

Are you satisfied with the resolution of your issue?
Yes
No

@MinasTyuru I just submitted a change that should help with this issue. Feel free to reopen if you encounter the segfault again.

Should col_impl get a similar treatment too? If col_ctx could be freed prematurely, then so could col_impl.

Thanks Ayush! I will test this out and see if it works.

Shimin, it looks to me like col_impl is a member of the col_ctx object and therefore does not need to be managed separately (it will be allocated and freed along with col_ctx). Oops, I meant col_params_, never mind.

@dubey It seems like the change fixes the col_ctx_ problem, however I'm still experiencing segfaults. I now observe that col_ctx_ is nonzero pointer, however col_ctx_->col_exec is a zero pointer, like so:

(gdb) bt
#0  0x00007fffbf8911a7 in tensorflow::NcclReducer::Run(std::function<void (tensorflow::Status const&)>) (this=0x7ffd9800f360, done=...)
    at external/org_tensorflow/tensorflow/core/kernels/collective_nccl_reducer.cc:185
#1  0x00007fffc83399e4 in tensorflow::BaseCollectiveExecutor::<lambda()>::operator() (__closure=<optimized out>)
    at external/org_tensorflow/tensorflow/core/common_runtime/base_collective_executor.cc:275
#2  std::_Function_handler<void(), tensorflow::BaseCollectiveExecutor::ExecuteAsync(tensorflow::OpKernelContext*, const tensorflow::CollectiveParams&, const string&, tensorflow::StatusCallback)::<lambda()> >::_M_invoke(const std::_Any_data &) (__functor=...)
    at external/gcc_7_4/usr/include/c++/7/bits/std_function.h:316
#3  0x00007fffbd5425c8 in std::function<void ()>::operator()() const (this=this@entry=0x7ffabcffce60)
    at external/gcc_7_4/usr/include/c++/7/bits/std_function.h:706
#4  0x00007fffc88674de in tensorflow::UnboundedWorkQueue::PooledThreadFunc (this=0x1ed20c40)
    at external/org_tensorflow/tensorflow/core/platform/default/unbounded_work_queue.cc:99
...
(gdb) p col_ctx_
$1 = std::shared_ptr<tensorflow::CollectiveContext> (use count 33194781, weak count -92074895) = {get() = 0x7ffd00555043}
(gdb) p col_params_
$2 = (const tensorflow::CollectiveParams *) 0x7fffc838ca46
     <std::_Function_handler<void(), tensorflow::(anonymous namespace)::ExecutorState::Process(tensorflow::(anonymous namespace)::ExecutorState::TaggedNode, tensorflow::int64)::<lambda()> >::_M_invoke(const std::_Any_data &)>
(gdb) p col_ctx_->col_exec
$3 = (tensorflow::CollectiveExecutor *) 0x0

I think that maybe CollectiveExecutor needs to be a shared pointer as well? And, as Shimin said, possibly col_impl also.

Also I think I don't have the required permissions to reopen this issue, so I'm just commenting on it.

Thanks for the update. Let me follow up internally.

Shimin implemented the change that he suggested with col_impl:

diff --git a/tensorflow/core/common_runtime/base_collective_executor.cc b/tensorflow/core/common_runtime/base_collective_executor.cc
index 0eebab9cf53..4d95059d018 100644
--- a/tensorflow/core/common_runtime/base_collective_executor.cc
+++ b/tensorflow/core/common_runtime/base_collective_executor.cc
@@ -246,20 +246,20 @@ void BaseCollectiveExecutor::ExecuteAsync(OpKernelContext* ctx,
                           col_params.is_source))
                             ? &ctx->input(0)
                             : nullptr;
-  CollectiveImplementationInterface* col_impl = nullptr;
-  Status status = CreateCollective(col_params, &col_impl);
+  CollectiveImplementationInterface* col_impl_ptr = nullptr;
+  Status status = CreateCollective(col_params, &col_impl_ptr);
   if (!status.ok()) {
     done_safe(status);
-    TF_DCHECK_EQ(nullptr, col_impl);
+    TF_DCHECK_EQ(nullptr, col_impl_ptr);
     return;
   }
+  auto col_impl = std::shared_ptr<CollectiveImplementationInterface>(col_impl_ptr);
   auto col_ctx = std::make_shared<CollectiveContext>(
       this, dev_mgr_, ctx, CtxParams(ctx), col_params, exec_key, step_id_,
       input, output);
   status = col_impl->InitializeCollectiveContext(col_ctx);
   if (!status.ok()) {
     done_safe(status);
-    delete col_impl;
     return;
   }
   // Run on an unbounded work queue that can handle blocking work so as to not
@@ -274,7 +274,6 @@ void BaseCollectiveExecutor::ExecuteAsync(OpKernelContext* ctx,
         profiler::TraceMeLevel::kInfo);
     col_impl->Run([col_impl, col_ctx, done_safe](const Status& s) {
       done_safe(s);
-      delete col_impl;
     });
   });
 }

I'm not sure what the relationship is between col_impl and the col_exec variable that seemed to be null previously, but it seems to fix the segfault. I'm doing some additional testing to be sure and will keep you posted.

Thanks, I was working on a similar change internally, this time with a unit test that can reproduce the issue. Yes please do keep me posted.

Apologies for the delay, it looks like everything works without segfaulting now.

Are you satisfied with the resolution of your issue?
Yes
No