Comments (6)
That's the correct thing to do. pmean
is a collective reduce (ie like MPI_AllReduce instead of MPI_Reduce, if you're familiar with MPI) and the result is sharded all devices. The logging call requires the data on the host, which transfers the array back, resulting in an array of length num_devices
, with elements, as you say, indentical due to the pmean
. The same thing is done in several places during the main training loop.
from ferminet.
That does make sense - might be just a small change from the internal version that we never tested. We'd be happy to look at a PR.
from ferminet.
Sorry, these crept in when removing a small internal dependency and as a typo respectively. Happy to take a look at a PR, equally happy to release an updated version early next week with fixes.
from ferminet.
Got it. Will work on a fix PR then
from ferminet.
By the way, for case 2 (the logging one),I got confused by the error message. It should be the case that constants.pmap
always returns an array (of the same values in this case) whose size is equal to num_devices
. It's just that when the array is of size 1, then logging can handle it automatically while if the size the greater than 1, it would complain. So I simply take the first element in my PR.
from ferminet.
Fixed in #17. Thanks for spotting this and sending the patch!
from ferminet.
Related Issues (20)
- Jax install - issue with correct version number HOT 1
- AttributeError: module 'jax.core' has no attribute 'extract_call_jaxpr' HOT 1
- Jax error running on A100 GPU (everything is okay on CPU) HOT 2
- unable to setup HOT 1
- The proper way to cite FermiNet repo HOT 1
- Ground State Energies HOT 2
- Question about pbc ewald part. HOT 2
- nan when training with 'adam' HOT 1
- About configs HOT 3
- Question About load Checkpoint HOT 1
- Evaluating logprob using batch_network in train HOT 1
- Issue on running pytest HOT 5
- Extension of PBC code to 1D HOT 7
- Something went wrong in RepeatedDenseBlock.update_curvature_matrix_estimate HOT 2
- Different results obtained from the paper for ch3nh2 HOT 2
- kfac_jax error when running H2 example script HOT 2
- Upstream breaking change in `kfac-jax`
- KeyError raised after burn-in MCMC steps HOT 1
- Logdet Bug Similar to e9f8c64 HOT 1
- Natural Excited States to nan
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.
from ferminet.