[CUDA] Set GPU device ID in threads

[shiyu1994]

This is to fix the issue proposed in #6018. As reported in the issue, currently illegal memory access may arise when using a gpu_device_id > 0 for cuda tree learner. This is because some CUDA memory is allocated in threads. Though the device ID is set in the main thread, it is not set in the newly created threads. The allocated CUDA memory in those threads may reside in GPU 0 by default, which is different with the gpu_device_id. Accessing to such memory on gpu_device_id may cause illegal memory access.

Here's one example,

LightGBM/src/io/cuda/cuda_column_data.cpp

Lines 120 to 137 in 20975ba

	#pragma omp parallel for schedule(static) num_threads(num_threads_)
	for (int column_index = 0; column_index < num_columns_; ++column_index) {
	OMP_LOOP_EX_BEGIN();
	const int8_t bit_type = column_bit_type[column_index];
	if (column_data[column_index] != nullptr) {
	// is dense column
	if (bit_type == 4) {
	column_bit_type_[column_index] = 8;
	InitOneColumnData<false, true, uint8_t>(column_data[column_index], nullptr, &data_by_column_[column_index]);
	} else if (bit_type == 8) {
	InitOneColumnData<false, false, uint8_t>(column_data[column_index], nullptr, &data_by_column_[column_index]);
	} else if (bit_type == 16) {
	InitOneColumnData<false, false, uint16_t>(column_data[column_index], nullptr, &data_by_column_[column_index]);
	} else if (bit_type == 32) {
	InitOneColumnData<false, false, uint32_t>(column_data[column_index], nullptr, &data_by_column_[column_index]);
	} else {
	Log::Fatal("Unknow column bit type %d", bit_type);
	}

[guolinke]

do we need to set this inside loop?

[shiyu1994]

Move outside. d4695ff

[jameslamb]

Thanks for the thorough explanation! It's really helpful for my understanding 😊

[jameslamb]

I just merged latest master into this, to get the fixes from #6032.

@shiyu1994 , could you go into the repo settings and check this "Always suggest updating pull request branches" box?

That would add a button that you can click on PRs to merge master in them, so there wouldn't be a need to pull and push locally. I've found that saves a bit of work in other projects I work on.

[shiyu1994]

could you go into the repo settings and check this "Always suggest updating pull request branches" box?

Done.

[github-actions]

This pull request has been automatically locked since there has not been any recent activity since it was closed. To start a new related discussion, open a new issue at https://github.com/microsoft/LightGBM/issues including a reference to this.