Skip to content

Rename batchsize to batchsize_per_gpu #475

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jul 10, 2025
Merged

Rename batchsize to batchsize_per_gpu #475

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b8d8fae
Rename batchsize to batchsize_per_gpu
sophie-xhonneux Jul 8, 2025
1537c86
Fix ruff stuff
sophie-xhonneux Jul 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions config/default_config.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,8 +70,8 @@ loss_fcts_val:
- "mse"
- 1.0

batch_size: 1
batch_size_validation: 1
batch_size_per_gpu: 1
batch_size_validation_per_gpu: 1

# training mode: "forecast" or "masking" (masked token modeling)
# for "masking" to train with auto-encoder mode, forecast_offset should be 0
Expand Down
12 changes: 8 additions & 4 deletions src/weathergen/model/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,8 +89,8 @@ def create(self, cf):
self.hp_nbours = torch.nn.Parameter(nbours, requires_grad=False)

# varlen index set for tokens
assert cf.batch_size == cf.batch_size_validation
bs = cf.batch_size
assert cf.batch_size_per_gpu == cf.batch_size_validation_per_gpu
bs = cf.batch_size_per_gpu
nqs = 9
s = [bs, self.num_healpix_cells, cf.ae_local_num_queries, cf.ae_global_dim_embed]
pad = torch.zeros(1, dtype=torch.int32)
Expand Down Expand Up @@ -467,7 +467,9 @@ def embed_cells(self, model_params, streams_data):

#########################################
def assimilate_local(self, model_params, tokens, cell_lens):
batch_size = self.cf.batch_size if self.training else self.cf.batch_size_validation
batch_size = (
self.cf.batch_size_per_gpu if self.training else self.cf.batch_size_validation_per_gpu
)

s = self.q_cells.shape
# print( f'{np.prod(np.array(tokens.shape))} :: {np.prod(np.array(s))}'
Expand Down Expand Up @@ -565,7 +567,9 @@ def forecast(self, model_params, tokens):
#########################################
def predict(self, model_params, fstep, tokens, streams_data, target_coords_idxs):
# fp32, i32 = torch.float32, torch.int32
batch_size = self.cf.batch_size if self.training else self.cf.batch_size_validation
batch_size = (
self.cf.batch_size_per_gpu if self.training else self.cf.batch_size_validation_per_gpu
)

s = [batch_size, self.num_healpix_cells, self.cf.ae_local_num_queries, tokens.shape[-1]]
tokens_stream = (tokens.reshape(s) + model_params.pe_global).flatten(0, 1)
Expand Down
32 changes: 16 additions & 16 deletions src/weathergen/train/trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,8 @@ def init(
):
self.cf = cf

assert cf.samples_per_epoch % cf.batch_size == 0
assert cf.samples_per_validation % cf.batch_size_validation == 0
assert cf.samples_per_epoch % cf.batch_size_per_gpu == 0
assert cf.samples_per_validation % cf.batch_size_validation_per_gpu == 0

self.mixed_precision_dtype = get_dtype(cf.attention_dtype)

Expand Down Expand Up @@ -87,7 +87,7 @@ def inference(self, cf, run_id_trained, epoch):
cf,
cf.start_date_val,
cf.end_date_val,
cf.batch_size_validation,
cf.batch_size_validation_per_gpu,
cf.samples_per_validation,
train_logger=self.train_logger,
stage=VAL,
Expand Down Expand Up @@ -141,7 +141,7 @@ def run(self, cf, run_id_contd=None, epoch_contd=None):
cf,
cf.start_date,
cf.end_date,
cf.batch_size,
cf.batch_size_per_gpu,
cf.samples_per_epoch,
train_logger=self.train_logger,
stage=TRAIN,
Expand All @@ -151,7 +151,7 @@ def run(self, cf, run_id_contd=None, epoch_contd=None):
cf,
cf.start_date_val,
cf.end_date_val,
cf.batch_size_validation,
cf.batch_size_validation_per_gpu,
cf.samples_per_validation,
train_logger=self.train_logger,
stage=VAL,
Expand Down Expand Up @@ -225,7 +225,7 @@ def run(self, cf, run_id_contd=None, epoch_contd=None):

# TODO: learning rate schedule
# https://www.cs.princeton.edu/~smalladi/blog/2024/01/22/SDEs-ScalingRules/
kappa = cf.batch_size * cf.num_ranks
kappa = cf.batch_size_per_gpu * cf.num_ranks
beta1 = max(0.5, 1.0 - kappa * (1.0 - 0.9))
beta2 = 1.0 - kappa * (1.0 - 0.999)
eps = 1e-08 / np.sqrt(kappa)
Expand All @@ -243,7 +243,7 @@ def run(self, cf, run_id_contd=None, epoch_contd=None):

# lr is updated after each batch so account for this
# TODO: conf should be read-only, do not modify the conf in flight
cf.lr_steps = int((len(self.dataset) * cf.num_epochs) / cf.batch_size)
cf.lr_steps = int((len(self.dataset) * cf.num_epochs) / cf.batch_size_per_gpu)

steps_decay = cf.lr_steps - cf.lr_steps_warmup - cf.lr_steps_cooldown
_logger.debug(f"steps_decay={steps_decay} lr_steps={cf.lr_steps}")
Expand All @@ -265,7 +265,7 @@ def run(self, cf, run_id_contd=None, epoch_contd=None):
_logger.warning(s)
self.lr_scheduler = LearningRateScheduler(
self.optimizer,
cf.batch_size,
cf.batch_size_per_gpu,
cf.num_ranks,
cf.lr_start,
cf.lr_max,
Expand Down Expand Up @@ -294,8 +294,8 @@ def run(self, cf, run_id_contd=None, epoch_contd=None):
epoch_base = int(self.cf.istep / len(self.data_loader))
else:
len_per_rank = (
len(self.dataset) // (self.num_ranks_original * cf.batch_size)
) * cf.batch_size
len(self.dataset) // (self.num_ranks_original * cf.batch_size_per_gpu)
) * cf.batch_size_per_gpu
epoch_base = int(
self.cf.istep / (min(len_per_rank, cf.samples_per_epoch) * self.num_ranks_original)
)
Expand Down Expand Up @@ -574,7 +574,7 @@ def train(self, epoch):
if bidx % self.checkpoint_freq == 0:
self.save_model(-1)

self.cf.istep += cf.batch_size
self.cf.istep += cf.batch_size_per_gpu

self.dataset.advance()

Expand Down Expand Up @@ -658,7 +658,7 @@ def validate(self, epoch):
self.losses_hist += [losses_all]
self.stddev_hist += [stddev_all]

pbar.update(self.cf.batch_size_validation)
pbar.update(self.cf.batch_size_validation_per_gpu)

losses_all = self.ddp_average(
torch.stack(self.losses_hist).to(torch.float64).nanmean(0)
Expand All @@ -678,7 +678,7 @@ def validate(self, epoch):
)

# add data to plain logger
samples = cf.istep * cf.batch_size * cf.num_ranks
samples = cf.istep * cf.batch_size_per_gpu * cf.num_ranks
self.train_logger.add_val(samples, losses_all, stddev_all)

if self.cf.rank == 0:
Expand Down Expand Up @@ -744,7 +744,7 @@ def log(self, bidx):
if bidx % log_interval == 0:
l_avg = self.ddp_average(torch.nanmean(torch.stack(self.losses_hist), axis=0))
stddev_avg = self.ddp_average(torch.nanmean(torch.stack(self.stddev_hist), axis=0))
samples = self.cf.istep * self.cf.batch_size * self.cf.num_ranks
samples = self.cf.istep * self.cf.batch_size_per_gpu * self.cf.num_ranks

if self.cf.rank == 0:
# logging
Expand Down Expand Up @@ -783,7 +783,7 @@ def log_terminal(self, bidx, epoch):
dt = time.time() - self.t_start
pstr = "{:03d} : {:05d}/{:05d} : {:06d} : loss = {:.4E} "
pstr += "(lr={:.2E}, s/sec={:.3f})"
len_dataset = len(self.data_loader) // self.cf.batch_size
len_dataset = len(self.data_loader) // self.cf.batch_size_per_gpu
print(
pstr.format(
epoch,
Expand All @@ -792,7 +792,7 @@ def log_terminal(self, bidx, epoch):
self.cf.istep,
np.nanmean(l_avg[0]),
self.lr_scheduler.get_lr(),
(self.print_freq * self.cf.batch_size) / dt,
(self.print_freq * self.cf.batch_size_per_gpu) / dt,
),
flush=True,
)
Expand Down