使用ScaLAPACK子程序PDGEMR2D出现内存不足错误

问题背景 #

使用ScaLAPACK的PDSYEV子程序计算大型哈密顿矩阵（N>15000）的全谱，计算过程表现正常。但将所有特征态合并到单个进程的全局矩阵时，N≤20000时pdgemr2d可以正常工作，N更大时触发内存不足错误：

cannot allocate memory, aborted xxmr2d: out of memory

集群配置：128GB内存、32CPU/64线程。

代码结构 #

.... 
   call MPI_INIT(info)
   call MPI_COMM_RANK(MPI_COMM_WORLD, myrank, info)
   call MPI_COMM_RANK(MPI_COMM_WORLD, nprocs, info)
   call MPI_COMM_SIZE(MPI_COMM_WORLD, nsize, info)
   nprow = 4
   npcol = 4
   block = 6
   call  blacs_get(-1,0,ictxt)
   call  blacs_gridinit(ictxt,"R",nprow, npcol)
   call  blacs_gridinfo(ictxt, nprow, npcol, myrow, mycol)
  ! print*, "myrank", myrank, myrow,mycol
   np = numroc(nst,block,myrow,0,nprow)
   nq = numroc(nst,block, mycol, 0,npcol)
   lld = max(1,np)
   call descinit(descAloc,nst,nst,block,block, 0,0, ictxt,np,info)
  ! print*, "INFO descinit A", info
  ! print*, "descAloc", descAloc
   call DESCINIT(descZ, nst, nst, block, block, 0, 0, ictxt,np, info)
   call DESCINIT(descZglobal, nst, nst, nst, nst, 0,0,ictxt,nst,info)
   call DESCINIT(descdumm, nst,nst,block,block,0,0,ictxt,np,info)
   !print*, "info descinit Z", info
   pi=acos(-1.d0) ! value of pi
   fi = 0.0/N_site
   phase =  dcmplx(cos(fi),sin(fi))
  
   !if  we  eigenvalues only 
   ! lworkk = 5*nst + max(3*block,block*(np+1))+1 
   !  if also eigenvectors 
    lworkk =  max( 1 + 6*nst + 2*np*nq, 3*nst + max( block*( np+1 ), 3*block)) + 2*nst
   allocate(Z(lld,lld), Aloc(lld,lld),workk(lworkk), wrk(np)) ! allocate memory to construct the full Hamiltonian matrix, j operator matrix for the calculated number of states in the Sz=0 sector 
      do local_i = 1 , np
        do local_j = 1 , nq
           ! Determine the global indicies (i,j) from the local indices   (local_i,local_j)
           i = indxl2g(local_i,block,myrow,0,nprow)
       j = indxl2g(local_j,block,mycol,0,npcol)
           ! Compute Hmatc(i,j) directly
           if (i == j) then 
           Aloc(local_i,local_j) = real(AD(i))
       else 
       Aloc(local_i, local_j) = 0.d0  ! Initialize to zero 
           do ih = row_ptr(i),row_ptr(i+1) - 1
          ij = col_ind(ih)
          isig = sign(1,col_ind_oper(ih))
              if (ij == j) then 
          if (isig == 1) then 
          Aloc(local_i,local_j) = Aloc(local_i,local_j) + 0.5*J_int
          else 
          Aloc(local_i,local_j) = Aloc(local_i,local_j) - 0.5*isig*J_int
              end if 
          end if 
           end do 
           end if 
        end do 
      end do 
   call pdsyev("v","u",nst,Aloc,1,1,descAloc,W,Z,1,1,descZ,workk,lworkk,info)
   if(myrank ==0) then 
   write(0,*) "info pdsyev", info 
   end if 
   if(myrank == 0) then
        write(0,*) "First and last eigenvalues :"
        write(0,"(6f5.2)") W(1),W(2),W(3),W(nst-2),W(nst-1),W(nst)
   write(0,*) "info pdsyev", info 
   end if     
   allocate(Zglobal(nst,nst)) ! global matrix to store all eigenvectors 
   call pdgemr2d(nst,nst, Z,1,1,descz, Zglobal, 1,1, descZglobal,ictxt)


   if(myrank == 0) then 
   deallocate(Zglobal)
   end if
    deallocate (Z)
    call blacs_gridexit(ictxt)
    call MPI_FINALIZE(info)

编译命令 #

mpiifx problem.f90 -qmkl=cluster
mpirun -np 16 ./a.out

错误信息 #

cannot allocate memory, aborted xxmr2d: out of memory

解决方案 #

1. 限制全局矩阵的分配范围 #

原代码中所有MPI进程都执行了allocate(Zglobal(nst,nst))，但只有rank=0进程需要存储全局特征向量矩阵。修改为仅在主进程分配：

if(myrank == 0) then
  allocate(Zglobal(nst,nst)) ! 仅主进程分配全局矩阵
end if

原因：当N=30000时，单个双精度Zglobal矩阵占用约7GB内存，16个进程同时分配会占用112GB，加上其他计算内存直接触发OOM。

2. 修正全局矩阵的描述符参数 #

原代码中descZglobal的块大小设为nst，导致ScaLAPACK尝试按整个矩阵块传输，瞬时内存压力极大。改为和descZ一致的块大小block，并仅在主进程初始化有效描述符：

! 初始化全局矩阵描述符
if(myrank == 0) then
  call DESCINIT(descZglobal, nst, nst, block, block, 0, 0, ictxt, nst, info)
else
  ! 其他进程初始化空描述符避免错误
  call DESCINIT(descZglobal, 0, 0, block, block, 0, 0, ictxt, 1, info)
end if

原因：按小块传输可以降低单进程的瞬时内存占用，避免一次性申请大量内存。

3. 分批次传输特征向量（针对超大N场景） #

如果N超过40000，即使主进程单独分配Zglobal也会占用过多内存，可以分批次传输特征向量：

• 每次传输K列特征向量到主进程
• 主进程处理完该批次后释放内存，再传输下一批
• 示例伪代码：

if(myrank == 0) then
  allocate(Zbatch(nst, K)) ! K为批次大小，比如1000
end if

do col_start = 1, nst, K
  col_end = min(col_start + K - 1, nst)
  call pdgemr2d(nst, col_end-col_start+1, Z, 1, col_start, descZ, &
                Zbatch, 1, 1, descZbatch, ictxt)
  ! 处理Zbatch中的特征向量
end do

if(myrank == 0) then
  deallocate(Zbatch)
end if

4. 内存预留与监控 #

• 计算主进程的内存需求：双精度矩阵内存= NN8字节，确保主进程有足够剩余内存（比如预留20GB以上给系统和其他计算）
• 运行前用free -h查看节点内存状态，避免其他进程占用内存

内容的提问来源于stack exchange，提问作者FortranZT