diff --git a/.gitignore b/.gitignore index a24c778..617fe48 100644 --- a/.gitignore +++ b/.gitignore @@ -34,7 +34,9 @@ examples/hBN/negf_output_k20/self_energy/self_energy_leadL.h5 examples/hBN/negf_output_k20/self_energy/self_energy_leadR.h5 examples/hBN/negf_output_k50/self_energy/* examples/hBN/negf_output_k70/self_energy/* -examples/CNT/output/* +examples/CNT/negf_profiling +examples/CNT/output* +examples/CNT/*long* examples/long_cnt/* CLAUDE* ai_docs/* diff --git a/Dockerfile b/Dockerfile index 1e020c2..244fc7a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -41,7 +41,7 @@ RUN \ conda run -n dpnegf pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cpu && \ # [2] Pull torch-scatter from the PyG wheel index and use --only-binary=torch-scatter to fully disable source builds. # If no matching wheel is found it fails immediately instead of spending ~10 minutes compiling a package that would crash at runtime. - conda run -n dpnegf pip install torch-scatter -f https://data.pyg.org/whl/torch-2.5.0+cpu.html --only-binary=torch-scatter && \ + conda run -n dpnegf pip install torch-scatter -f https://data.pyg.org/whl/torch-2.5.1+cpu.html --only-binary=torch-scatter && \ # [3] Guard the local-repo installs with the CPU index so hidden dependencies can't replace the CPU torch with a CUDA build. conda run -n dpnegf pip install ./DeePTB torch==2.5.1 --extra-index-url https://download.pytorch.org/whl/cpu && \ conda run -n dpnegf pip install ./ torch==2.5.1 --extra-index-url https://download.pytorch.org/whl/cpu && \ diff --git a/dpnegf/negf/recursive_green_cal.py b/dpnegf/negf/recursive_green_cal.py index 3cefcca..9d62337 100644 --- a/dpnegf/negf/recursive_green_cal.py +++ b/dpnegf/negf/recursive_green_cal.py @@ -176,10 +176,17 @@ def recursive_gf_cal(energy, mat_l_list, mat_d_list, mat_u_list, # In-place: mat_d_list is a fresh tensor (wrapper's `* 1.` copy on D), # so we can fuse the energy shift without the e_bcast*sd transient. mat_d_list[jj].addcmul_(sd[jj], e_bcast, value=-1) + # sd[jj] is dead after this — it's only read here in the non-uniform + # kernel. Drop the wrapper-side view so the addcmul_ transient slab + # can be coalesced by the caching allocator inside the loop instead + # of waiting for the wrapper's `sd_b` Python name to leave scope. + sd[jj] = None for jj in range(len(mat_l_list)): mat_l_list[jj] = mat_l_list[jj] - e_bcast * sl[jj] + sl[jj] = None for jj in range(len(mat_u_list)): mat_u_list[jj] = mat_u_list[jj] - e_bcast * su[jj] + su[jj] = None num_of_matrices = len(mat_d_list) mat_shapes = [item.shape for item in mat_d_list] # [B, n_q, n_q] @@ -454,6 +461,13 @@ def _to_batch(t): Sd = torch.stack(sd_b, dim=0) # [K, B, n, n] Sl = torch.stack(sl_b, dim=0) # [K-1, B, n, n] Su = torch.stack(su_b, dim=0) # [K-1, B, n, n] + # torch.stack on the wrapper's `*1.` D copies and the L/U/sd/sl/su + # expanded views produces six owned 4-D tensors. The wrapper-side + # lists are dead from here on — drop them now so the per-slot + # `[B, n_q, n_q]` storage (≈ K × B × n² × 16 B for D) can be freed + # before the kernel allocates gr_left/grl/gru. + del temp_mat_d_list, temp_mat_l_list, temp_mat_u_list + del sd_b, sl_b, su_b ans = recursive_gf_cal(shift_energy, L, D, U, Sd, Su, Sl, s_in=s_in_b, s_out=s_out_b, eta=eta, need_lesser=need_lesser, @@ -470,6 +484,12 @@ def _to_batch(t): need_gr_lc=need_gr_lc, stacked=False, keep_gr_left=keep_gr_left) + # Non-uniform kernel consumed the lists by reference and nulled + # individual slots as it went. Drop the wrapper-side names so the + # Python list objects (and any straggler refs) are gone before + # _squeeze_ans/return. + temp_mat_d_list = temp_mat_l_list = temp_mat_u_list = None + sd_b = sl_b = su_b = None if squeezed: ans = _squeeze_ans(ans) diff --git a/dpnegf/runner/NEGF.py b/dpnegf/runner/NEGF.py index 77520bd..d18a134 100644 --- a/dpnegf/runner/NEGF.py +++ b/dpnegf/runner/NEGF.py @@ -607,9 +607,12 @@ def _auto_chunk_size(self, n_grid): Per-energy peak (post per-slot-release, complex128) approximated as bytes_per_E ~= C * K * n_max**2 * 16 with C bundling the live tensors in the worst backward-sweep slot - (grd full + grl + gru full + decaying gr_left tail + gU + transients). - C=10 with a 0.7x free-memory budget; deliberately conservative because - without expandable_segments the allocator can't defragment on demand. + (grd full + grl + gru full + decaying gr_left tail + gU + transients + + the K-resident H/S diagonal & off-diagonal lists that survive across + chunks). C=14 with a 0.5x free-memory budget; deliberately conservative + because without expandable_segments the allocator can't defragment on + demand, and on real workloads (CNT10/long6) the previous 10x / 0.7 + combination still picked a chunk that OOM'd on a 15.77 GiB V100. """ rgf_dev = self.rgf_device if not (isinstance(rgf_dev, torch.device) and rgf_dev.type == "cuda"): @@ -620,10 +623,10 @@ def _auto_chunk_size(self, n_grid): K = len(self.deviceprop.hd) except Exception: return n_grid - per_e = 10 * K * (n_max ** 2) * 16 + per_e = 14 * K * (n_max ** 2) * 16 if per_e <= 0: return n_grid - b = max(1, min(n_grid, int(0.7 * free_bytes) // per_e)) + b = max(1, min(n_grid, int(0.5 * free_bytes) // per_e)) log.info( f"auto e_batch_size={b} (free={free_bytes/2**30:.2f} GiB, " f"per_E~={per_e/2**20:.1f} MiB, K={K}, n_max={n_max})" @@ -751,8 +754,19 @@ def negf_compute(self,scf_require=False,Vbias=None): # Non-SCF: solve a whole chunk of energies in one batched recursive_gf call. if self.e_batch_size is not None: chunk = self.e_batch_size + # The user-supplied value bypasses the auto-budget. + rgf_dev = self.rgf_device + if isinstance(rgf_dev, torch.device) and rgf_dev.type == "cuda": + cap = self._auto_chunk_size(len(self.uni_grid)) + if chunk > cap: + log.warning( + f"user e_batch_size={chunk} exceeds the " + f"CUDA auto-cap={cap} on {rgf_dev}; " + ) else: chunk = self._auto_chunk_size(len(self.uni_grid)) + log.info(f"Using e_batch_size={chunk} for energy loop with {len(self.uni_grid)} points") + for e_chunk in torch.split(self.uni_grid, chunk): e_batch_size = len(e_chunk) log.info( diff --git a/ut.sh b/ut.sh index 13761eb..c326aa9 100644 --- a/ut.sh +++ b/ut.sh @@ -19,7 +19,7 @@ echo "--- Installing/updating package from PR in editable mode ---" # Docker image. # 2. `pytest ./tests/`: After the package is installed, we run the tests. -conda run -n dpnegf bash -c "pip install -e . 'torch==2.1.1' --extra-index-url https://download.pytorch.org/whl/cpu -f https://data.pyg.org/whl/torch-2.1.1+cpu.html --only-binary=torch-scatter && pytest dpnegf/tests/" +conda run -n dpnegf bash -c "pip install -e . 'torch==2.5.1' --extra-index-url https://download.pytorch.org/whl/cpu -f https://data.pyg.org/whl/torch-2.5.1+cpu.html --only-binary=torch-scatter && pytest dpnegf/tests/" echo "--- Unit Tests Passed Successfully ---"