lberrada
commented
11 months ago Hi, sorry about the slow reply.
Here are some answers that will hopefully clarify things:
- you're correct that
--jaxline_mode=train_eval_multithreadedis required to run evaluation, otherwise only training is applied. This is an artefact of the JAXline framework that is currently used. global_stepis a counter that gets incremented at each step, whether computing a gradient accumulation step for the virtual batch-size or actually performing an update. The name is unfortunate, but this is also something that JAXline enforces, and that we cannot easily modify.update_stepis indeed incremented whenever a model update is performed.- the logs are not shown for each step, which is why you observe the strange "skipping" pattern. If you set
log_train_data_interval=1here, you should see the logs at everyglobal_stepand the pattern should become clearer. - Another subtlety about
dp_epsilonis that we actually pre-compute values and cache them on a regular grid of the steps. This avoids running slow accounting operations on CPU at each training step, which would slow things down. Sodp_epsilonis actually a piecewise constant approximation as a function of the steps, where each cached value within a constant piece should be an upper bound on the "exact"dp_epsilon.
We're aware that global_step vs update_step is confusing, and that the eval setup is not ideal; these issues have been stemming from constraints of JAXline. The good news though is these confusing bits mentioned so far (eval being difficult to run, global step vs update step) will go away in a forthcoming version that I hope to release soon!
SabrinaMokhtari
Hello there!
I'm currently working on the Image Classification experiments.
When trying to run the code for CIFAR fine-tuning or training from scratch on CIFAR-10, using --jaxline_mode=train_eval_multithreaded triggers this error:
However, removing this mode seems to prevent the code from going through the evaluation step, or at least, the logs don't capture it. All the logs seem to pertain to training, following the 'train/parameter' format as illustrated below.
On another note, I'm curious about the batch size. The default hyperparameters suggest a training batch size of 4096 and a per_device_per_step of 64. However, according to the logs, the data_seen at each step is 6400 which seems to align with global steps per_device_per_step (100 64) but not with the original batch size. Another confusing factor is the train/update_step, which appears to lack a specific order (it goes from 1 to 3, 4, 6, 7, ...). Understanding the relationship between batch size (4096), per_device_per_step (64), the number of global steps (100), update_step, and epochs has become a bit of a puzzle.
Additionally, I've noticed that the dp_epsilon value changes only after 7 log prints. I presume this behavior originates from taking virtual steps and working with a virtual batch size. However, making sense of this and the various numbers mentioned earlier is proving to be quite a challenge. Any insights or clarifications would be highly appreciated!
Thanks!