如何写好一个分布式程序,可以同时支持单节点多核/多节点/多gpu?

julia初学者。尝试将一个fortran程序用julia改写,使其能在单节点多核/多节点/多gpu等各种环境下执行。

1.应该如何写比较好?是否有什么库,可以支持多种array(比如sharedarray,cudarray,distributed array等;还有双精度和单精度),从而能让我在改写时尽量复用代码?
2.原本的fortran程序用的是mpi,我改写的话,用什么会比较好?依旧用mpi还是用distributed array?使用mpi和distributed array在多节点上运行,在效率或其他方面有什么区别?
3.有没有类似pytorch中amp(自动混合精度)的库,使我的程序可以非常方便的在乘法时自动使用较低精度?(比如,若双精度数据,在乘法时自动转为单精度;单精度乘法自动转为半精度)

第一点要求可能有点高,毕竟CuArray计算的编写模式和普通Array还是区别很大,ParallelStencil.jl的一些功能也许可以用作参考。
第二点会很取决于计算的规模,英文discourse里有人作过比较,在大规模计算时MPI.jl比Distributed有一些性能优势。

根据我几年前的对于一维扩散问题的测试结果,Julia中的distributed array的性能比shared array差,且二者性能表现都不是很理想。

Problem size Serial [s] SArray [s] DArray [s]
100 0.0 0.19 2.05
1000 0.002 0.20 13.03
10000 0.019 0.50 125.00
100000 0.197 0.97 ???
1000000 2.264 4.89 ???

当然可能是我没有找到正确的使用方法,但大规模并行改写还是推荐从MPI开始。

你的多平台并行需求目前我没有发现在Julia社区中有以此为卖点的库。ParallelStencil.jl可能还是以GPU为出发点的,多核多线程的并行大概率是没有仔细优化的。精度自动转换部分,这个需求感觉是机器学习引发的,如果你有找到相应功能的底层库可以跟大家分享。

大家好。经过我一天的初步研究,查阅了几个在大气、海洋方面的大型项目,得出了以下初步结论:

  1. 几乎所有这个领域的大型julia项目都是使用mpi进行并行;
  2. 几乎所有这个领域的大型julia项目都是使用两套代码,分别应对cpu/gpu两套环境;
  3. 相对而言,c++和python都有跨device的基础数组库,而julia上并没有找到。
  4. 我在fortran中找到了一个手动混合精度的库,在一个海洋模式上获得了较好的效果,说明混合进度似乎确实能有效提高计算速度。在python上,也是使用手动混合进度较多,自动混合进度只找到了pytorch的AMP。
2 个赞

我很好奇这种混合精度对于计算结果尤其是PDE的精度影响。你在哪里有看到分析吗?

[1] Prims, O. T. , Acosta, M. C. , Moore, A. M. , Castrillo, M. , & Doblas-Reyes, F. J. . (2019). How to use mixed precision in ocean models: exploring a potential reduction of numerical precision in nemo 4.0 and roms 3.6. Geoscientific Model Development, 12(7), 3135-3148.

我是看到这个文章里的数据。它用了一个fortran库,支持了混合精度,然后改写了一个海洋模式(全部在cpu上),最后的结论是精度没太大变化,速度和内存使用都减少了

2 个赞

c++和python都有跨device的基础数组库 是啥啊? Processors · Dagger.jl 这个是不是julia目前最接近你说的?

c++ 上是kokkos,python可以用jax; dagger.jl对标的应该是python的dask,和我说的没啥关系。。。

你这个试试Elixir + Nx:joy:
要是我的话,可能得用Julia手撸一个自己用的工具。Threads + Distributed

这让我想起来 GitHub - CliMA/Oceananigans.jl: 🌊 Julia software for fast, friendly, flexible, ocean-flavored fluid dynamics on CPUs and GPUs