diff --git a/graphphysics/dataset/dataset.py b/graphphysics/dataset/dataset.py index 22f588e..fb6cb69 100644 --- a/graphphysics/dataset/dataset.py +++ b/graphphysics/dataset/dataset.py @@ -26,6 +26,7 @@ def __init__( add_edge_features: bool = True, use_previous_data: bool = False, world_pos_parameters: Optional[dict] = None, + target_same_frame: bool = False ): with open(meta_path, "r") as fp: meta = json.load(fp) @@ -50,7 +51,7 @@ def __init__( self.new_edges_ratio = new_edges_ratio self.add_edge_features = add_edge_features self.use_previous_data = use_previous_data - + self.target_same_frame = target_same_frame self.world_pos_index_start = None self.world_pos_index_end = None if world_pos_parameters is not None: @@ -77,12 +78,20 @@ def get_traj_frame(self, index: int) -> Tuple[int, int]: Returns: Tuple[int, int]: A tuple containing the trajectory number and the frame number within that trajectory. """ - traj = index // (self.trajectory_length - 1) - frame = index % (self.trajectory_length - 1) + int(self.use_previous_data) + if (self.target_same_frame): + traj = index // self.trajectory_length + frame = index % self.trajectory_length + else: + traj = index // (self.trajectory_length - 1) + frame = index % (self.trajectory_length - 1) + int(self.use_previous_data) return traj, frame def __len__(self) -> int: - return self.size_dataset * (self.trajectory_length - 1) + self.target_same_frame=True # OJO !!!!!!!!!!!!!!!!!!!!!!!! + if (self.target_same_frame): + return self.size_dataset*self.trajectory_length + else: + return self.size_dataset * (self.trajectory_length - 1) @abstractmethod def __getitem__(self, index: int) -> Data: diff --git a/graphphysics/dataset/xdmf_dataset.py b/graphphysics/dataset/xdmf_dataset.py index 3fd6e25..512ce62 100644 --- a/graphphysics/dataset/xdmf_dataset.py +++ b/graphphysics/dataset/xdmf_dataset.py @@ -24,6 +24,7 @@ def __init__( add_edge_features: bool = True, use_previous_data: bool = False, switch_to_val: bool = False, + target_same_frame: bool = False, random_prev: int = 1, # If we use previous data, we will fetch one previous frame between [-1, -random_prev] random_next: int = 1, # The target will be the frame : t + [1, random_next] ): @@ -35,6 +36,7 @@ def __init__( new_edges_ratio=new_edges_ratio, add_edge_features=add_edge_features, use_previous_data=use_previous_data, + target_same_frame=target_same_frame, ) self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -56,12 +58,13 @@ def __init__( self.xdmf_folder = xdmf_folder self.meta_path = meta_path + self.target_same_frame = target_same_frame # Get list of XDMF files in the folder self.file_paths: List[str] = [ os.path.join(xdmf_folder, f) for f in os.listdir(xdmf_folder) - if os.path.isfile(os.path.join(xdmf_folder, f)) and f.endswith(".xdmf") + if os.path.isfile(os.path.join(xdmf_folder, f)) and (f.endswith(".xdmf") or f.endswith(".xmf")) ] self._size_dataset: int = len(self.file_paths) @@ -94,32 +97,57 @@ def __getitem__(self, index: int) -> Union[Data, Tuple[Data, torch.Tensor]]: _previous_data_index = random.randint(1, self.random_prev) # Read XDMF file - with meshio.xdmf.TimeSeriesReader(xdmf_file) as reader: - num_steps = reader.num_steps - - if frame - _previous_data_index < 0: - _previous_data_index = 1 - if frame + _target_data_index > num_steps - 1: - _target_data_index = 1 - - if frame >= num_steps - 1: - raise IndexError( - f"Frame index {frame} out of bounds for trajectory {traj_index} with {num_steps} frames." - ) - - points, cells = reader.read_points_cells() - time, point_data, _ = reader.read_data(frame) - _, target_point_data, _ = reader.read_data(frame + _target_data_index) - - if self.use_previous_data: - _, previous_data, _ = reader.read_data(frame - _previous_data_index) - + # --- ADD --- + # Option without time serie + if self.target_same_frame: + mesh = meshio.read(xdmf_file) + points, cells = mesh.points, mesh.cells + point_data = dict(mesh.point_data) + # --- AJOUT : inclure mesh_pos et wall_mask --- + point_data["mesh_pos"] = mesh.points.astype( + self.meta["features"]["mesh_pos"]["dtype"] + ) + if "wall_mask" in mesh.point_data.keys(): + point_data["wall_mask"] = np.array( + mesh.point_data["wall_mask"] + ).astype(self.meta["features"]["wall_mask"]["dtype"]) + # --- FIN AJOUT --- + target_point_data = point_data # même frame = même cible + previous_data = None + num_steps = 1 + target_frame = 1 + + # --- END ADD --- + else: + with meshio.xdmf.TimeSeriesReader(xdmf_file) as reader: + num_steps = reader.num_steps + + if frame - _previous_data_index < 0: + _previous_data_index = 1 + if frame + _target_data_index > num_steps - 1: + _target_data_index = 1 + + if frame >= num_steps - 1 and (not self.target_same_frame): + raise IndexError( + f"Frame index {frame} out of bounds for trajectory {traj_index} with {num_steps} frames.") + + points, cells = reader.read_points_cells() + time, point_data, _ = reader.read_data(frame) + target_frame = frame + 1 + if self.target_same_frame: target_frame = frame + _, target_point_data, _ = reader.read_data(target_frame) + + if self.use_previous_data: + _, previous_data, _ = reader.read_data(frame - _previous_data_index) + # Prepare the mesh data mesh = meshio.Mesh(points, cells, point_data=point_data) # Get faces or cells if "triangle" in mesh.cells_dict: cells = mesh.cells_dict["triangle"] + elif "line" in mesh.cells_dict: + cells = mesh.cells_dict["line"] elif "tetra" in mesh.cells_dict: cells = torch.tensor(mesh.cells_dict["tetra"], dtype=torch.long) else: @@ -128,18 +156,46 @@ def __getitem__(self, index: int) -> Union[Data, Tuple[Data, torch.Tensor]]: ) # Process point data and target data - point_data = { - k: np.array(mesh.point_data[k]).astype(self.meta["features"][k]["dtype"]) - for k in self.meta["features"] - if k in mesh.point_data.keys() - } - - target_data = { - k: np.array(target_point_data[k]).astype(self.meta["features"][k]["dtype"]) - for k in self.meta["features"] - if k in target_point_data.keys() - and self.meta["features"][k]["type"] == "dynamic" - } + if self.target_same_frame: + selected_features = ["mesh_pos", "wall_mask"] + point_data = { + k: np.array(mesh.point_data[k]).astype(self.meta["features"][k]["dtype"]) + for k in self.meta["features"] + if k in mesh.point_data.keys() and k in selected_features + } + + target_data = { + k: np.array(target_point_data[k]).astype(self.meta["features"][k]["dtype"]) + for k in self.meta["features"] + if k in target_point_data.keys() + and k == "Velocity" + } + else: + point_data = { + k: np.array(mesh.point_data[k]).astype(self.meta["features"][k]["dtype"]) + for k in self.meta["features"] + if k in mesh.point_data.keys() + } + + target_data = { + k: np.array(target_point_data[k]).astype(self.meta["features"][k]["dtype"]) + for k in self.meta["features"] + if k in target_point_data.keys() + and self.meta["features"][k]["type"] == "dynamic" + } + # --- DEBUG PRINT ESSENTIEL --- + #print('-----DEBUG--------') + #def summarize_data(name, data_dict): + # print(f"\n{name} summary:") + # for k, v in data_dict.items(): + # print(f" {k:<20} shape={v.shape}, dtype={v.dtype}") + # print(f"Total {len(data_dict)} features.\n") + + #summarize_data("point_data", point_data) + #summarize_data("target_data", target_data) + #print('-----------------') + # --- FIN DEBUG --- + #print(f"[DEBUG] GPU Mem Alloc: {torch.cuda.memory_allocated()/1024**2:.1f} MB, Reserved: {torch.cuda.memory_reserved()/1024**2:.1f} MB") def _reshape_array(a: dict): for k, v in a.items(): @@ -150,17 +206,28 @@ def _reshape_array(a: dict): _reshape_array(target_data) # Create graph from mesh data - graph = meshdata_to_graph( - points=points.astype(np.float32), - cells=cells, - point_data=point_data, - time=time, - target=target_data, - id=mesh_id, - ) + if self.target_same_frame: + graph = meshdata_to_graph( + points=points.astype(np.float32), + cells=cells, + point_data=point_data, + time=0, + target=target_data, + id=mesh_id, + ) + else: + graph = meshdata_to_graph( + points=points.astype(np.float32), + cells=cells, + point_data=point_data, + time=time, + target=target_data, + id=mesh_id, + ) + # TODO: add target_dt and previous_dt as features per node. graph.target_dt = _target_data_index * self.dt - + self.use_previous_data = False #### I force use_previous_data to be false if self.use_previous_data: previous = { k: np.array(previous_data[k]).astype(self.meta["features"][k]["dtype"]) @@ -174,7 +241,7 @@ def _reshape_array(a: dict): graph = graph.to(self.device) - graph = self._apply_preprocessing(graph) + #graph = self._apply_preprocessing(graph) #### I remove preprocess graph = self._apply_k_hop(graph, traj_index) graph = self._may_remove_edges_attr(graph) graph = self._add_random_edges(graph) diff --git a/graphphysics/models/layers.py b/graphphysics/models/layers.py index 8a27abe..42487fb 100644 --- a/graphphysics/models/layers.py +++ b/graphphysics/models/layers.py @@ -5,6 +5,59 @@ import torch.nn as nn from torch_geometric.nn import MessagePassing +# --- Helper SpMM: force l'op en FP32 pour éviter mismatch sous BF16 --- +import torch +from contextlib import contextmanager + +# Activation checkpoint +from torch.utils.checkpoint import checkpoint + +@contextmanager +def _no_autocast_cuda(): + # évite que autocast BF16 ré-intercepte SpMM + if torch.is_autocast_enabled(): + with torch.cuda.amp.autocast(enabled=False): + yield + else: + yield + +def spmm_fp32(adj, x): + """ + Effectue Y = adj @ x en FP32 (hors autocast), puis cast Y vers x.dtype. + Pourquoi: DGL SpMM exige même dtype entre valeurs(A) et X; BF16 n'est pas toujours supporté. + """ + x32 = x.float() + # Certaines versions DGL ont .astype ; sinon, le backend convertira côté C. + try: + adj32 = adj.astype(torch.float32) # DGL SparseMatrix + except AttributeError: + adj32 = adj # fallback: beaucoup de builds acceptent A en fp32 par défaut + + with _no_autocast_cuda(): + y32 = adj32 @ x32 # déclenche torch.ops.dgl_sparse.spmm(...) + return y32.to(x.dtype) + +# --- Helpers DGL sparse en FP32 (BF16-safe) --- +def bsddmm_fp32(mask, q, kT): + q32, kT32 = q.float(), kT.float() + try: + mask32 = mask.astype(torch.float32) + except AttributeError: + mask32 = mask + with _no_autocast_cuda(): + out = dglsp.bsddmm(mask32, q32, kT32) + return out # SparseMatrix (valeurs fp32) + +def bspmm_fp32(attn, v, out_dtype): + v32 = v.float() + try: + attn32 = attn.astype(torch.float32) + except AttributeError: + attn32 = attn + with _no_autocast_cuda(): + y32 = dglsp.bspmm(attn32, v32) + return y32.to(out_dtype) + try: import dgl.sparse as dglsp from dgl.sparse import SparseMatrix @@ -343,7 +396,8 @@ def scaled_query_key_softmax( q = q / scaling_factor if att_mask is not None and HAS_DGL_SPARSE: - attn = dglsp.bsddmm(att_mask, q, k.transpose(1, 0)) + #attn = dglsp.bsddmm(att_mask, q, k.transpose(1, 0)) + attn = bsddmm_fp32(att_mask, q, k.transpose(1, 0)).softmax() attn = attn.softmax() else: attn = q @ k.transpose(-2, -1) @@ -378,7 +432,8 @@ def scaled_dot_product_attention( # Compute the output if att_mask is not None and HAS_DGL_SPARSE: - y = dglsp.bspmm(attn, v) + #y = dglsp.bspmm(attn, v) + y = bspmm_fp32(attn, v, v.dtype) else: y = attn @ v @@ -520,12 +575,14 @@ def __init__( self.activation = activation_layer() self.norm1, self.norm2 = RMSNorm(output_dim), RMSNorm(output_dim) self.gated_mlp = build_gated_mlp( - in_size=output_dim, hidden_size=output_dim, out_size=output_dim + in_size=output_dim, hidden_size=output_dim, out_size=output_dim, + expansion_factor=2 # <-- change clé : 3 -> 2 ) self.use_adjacency = HAS_DGL_SPARSE + self.use_activation_checkpointing = False # togglé depuis l'extérieur - def forward( + '''def forward( self, x: torch.Tensor, adj, return_attention: bool = False ) -> torch.Tensor: """ @@ -556,7 +613,38 @@ def forward( return x, attn else: return x + ''' + def forward(self, x: torch.Tensor, adj, return_attention: bool = False) -> torch.Tensor: + if not self.use_adjacency: + adj = None + if return_attention: + # En validation/diag, on évite le checkpoint pour récupérer les poids d’attention + x_, attn = self.attention(self.norm1(x), adj, return_attention=True) + x = x + x_ + x = x + self.gated_mlp(self.norm2(x)) + return x, attn + + # --- En TRAIN: checkpoint fin pour réduire les activations gardées en mémoire --- + + # 1) Attention + def _attn_only(tensor_x): + # adj est capturé par la closure (non Tensor) => OK pour checkpoint + return self.attention(self.norm1(tensor_x), adj, return_attention=False) + + # 2) MLP + def _mlp_only(tensor_x): + return self.gated_mlp(self.norm2(tensor_x)) + + if self.training: + # use_reentrant=False consomme moins de mémoire avec PyTorch ≥1.12/2.0 + x = x + checkpoint(_attn_only, x, use_reentrant=False) + x = x + checkpoint(_mlp_only, x, use_reentrant=False) + else: + x = x + self.attention(self.norm1(x), adj, return_attention=False) + x = x + self.gated_mlp(self.norm2(x)) + + return x class GraphNetBlock(MessagePassing): """ diff --git a/graphphysics/train.py b/graphphysics/train.py index 28b82fd..cecd801 100644 --- a/graphphysics/train.py +++ b/graphphysics/train.py @@ -21,6 +21,12 @@ ) from graphphysics.utils.progressbar import ColabProgressBar +import socket +import torch.distributed as dist +from lightning.pytorch.strategies import DDPStrategy +from datetime import timedelta # si tu ne l’as pas déjà + + warnings.filterwarnings( "ignore", ".*Trying to infer the `batch_size` from an ambiguous collection.*" ) @@ -59,6 +65,24 @@ "training_parameters_path", None, "Path to the training parameters JSON file" ) +def print_dist_info(stage: str): + """Affiche les variables de distribution pour debug (appel court).""" + hostname = socket.gethostname() + env = os.environ + rank = int(env.get("RANK", env.get("SLURM_PROCID", -1))) + local_rank = int(env.get("LOCAL_RANK", env.get("SLURM_LOCALID", -1))) + node_rank = int(env.get("NODE_RANK", env.get("SLURM_NODEID", -1))) + world_size = int(env.get("WORLD_SIZE", env.get("SLURM_NTASKS", -1))) + initialized = dist.is_available() and dist.is_initialized() + num_gpus = torch.cuda.device_count() + curr = torch.cuda.current_device() if torch.cuda.is_available() and num_gpus > 0 else -1 + gpu_name = torch.cuda.get_device_name(curr) if curr != -1 else "CPU" + print( + f"[{stage}] host={hostname} rank={rank} local_rank={local_rank} node_rank={node_rank} " + f"world_size={world_size} dist_init={initialized} gpus_on_node={num_gpus} " + f"current_device={curr} device_name={gpu_name}", + flush=True, + ) def main(argv): del argv @@ -78,8 +102,9 @@ def main(argv): return device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + preproc_device = torch.device("cpu") - wandb_project_name = FLAGS.project_name + #wandb_project_name = FLAGS.project_name num_epochs = FLAGS.num_epochs initial_lr = FLAGS.init_lr batch_size = FLAGS.batch_size @@ -96,10 +121,14 @@ def main(argv): seed_everything(FLAGS.seed, workers=True) + print_dist_info("startup") + + # Build preprocessing function train_preprocessing = get_preprocessing( param=parameters, - device=device, + #device=device, + device=preproc_device, use_edge_feature=use_edge_feature, extra_node_features=build_features, ) @@ -114,7 +143,8 @@ def main(argv): val_preprocessing = get_preprocessing( param=parameters, - device=device, + #device=device, + device=preproc_device, use_edge_feature=use_edge_feature, remove_noise=True, extra_node_features=build_features, @@ -128,14 +158,44 @@ def main(argv): switch_to_val=True, ) + print("---- DEBUG ----") + print("TRAIN Dataset type:", type(train_dataset)) + print("Number of XDMF files:", train_dataset.size_dataset) + print("Trajectory length:", train_dataset.trajectory_length) + print("Computed len(train_dataset):", len(train_dataset)) + print("----------------") + print("VAL Dataset type:", type(val_dataset)) + print("Number of XDMF files:", val_dataset.size_dataset) + print("Trajectory length:", val_dataset.trajectory_length) + print("Computed len(val_dataset):", len(val_dataset)) + print("----------------") + num_workers = get_num_workers(param=parameters, default_num_workers=num_workers) + rank_env = int(os.environ.get("RANK", os.environ.get("SLURM_PROCID", 0))) + if rank_env == 0: + print("---- DEBUG (rank 0) ----") + print("TRAIN Dataset type:", type(train_dataset)) + print("Number of XDMF files:", train_dataset.size_dataset) + print("Trajectory length:", train_dataset.trajectory_length) + print("Computed len(train_dataset):", len(train_dataset)) + print("----------------") + print("VAL Dataset type:", type(val_dataset)) + print("Number of XDMF files:", val_dataset.size_dataset) + print("Trajectory length:", val_dataset.trajectory_length) + print("Computed len(val_dataset):", len(val_dataset)) + print("----------------") + else: + # Court message pour confirmer la présence des autres ranks + print(f"[rank {rank_env}] len(train)={len(train_dataset)} len(val)={len(val_dataset)}", flush=True) + train_dataloader_kwargs = { "dataset": train_dataset, "shuffle": True, "batch_size": batch_size, "num_workers": num_workers, "exclude_keys": ["tetra"], + "drop_last": True, } valid_dataloader_kwargs = { @@ -205,6 +265,7 @@ def main(argv): ) # Initialize WandbLogger + ''' if resume_training: wandb_run = wandb.init( project=wandb_project_name, id=lightning_module.wandb_run_id, resume="allow" @@ -214,6 +275,7 @@ def main(argv): wandb_logger = WandbLogger(experiment=wandb_run) lightning_module.wandb_run_id = wandb_logger.experiment.id + ''' if model_save_name is not None: checkpoint_callback = ModelCheckpoint( dirpath="checkpoints/", filename=model_save_name @@ -221,7 +283,7 @@ def main(argv): else: checkpoint_callback = ModelCheckpoint(dirpath="checkpoints") lr_monitor = LearningRateMonitor(logging_interval="step") - + ''' wandb_logger.experiment.config.update( { "architecture": parameters["model"]["type"], @@ -232,22 +294,55 @@ def main(argv): "batch_size": batch_size, } ) - + ''' + ''' # Configure Trainer trainer = Trainer( accelerator="gpu" if torch.cuda.is_available() else "cpu", devices=1, max_epochs=num_epochs, - logger=wandb_logger, + #logger=wandb_logger, callbacks=[ ColabProgressBar(), checkpoint_callback, - LogPyVistaPredictionsCallback(dataset=val_dataset, indices=[1, 2, 3]), + #LogPyVistaPredictionsCallback(dataset=val_dataset, indices=[1, 2, 3]), lr_monitor, ], log_every_n_steps=100, gradient_clip_val=1.0, ) + ''' + # === Trainer DDP === + num_nodes = int(os.environ.get("SLURM_NNODES", 1)) + devices = torch.cuda.device_count() if torch.cuda.is_available() else 0 + if devices == 0: + raise RuntimeError("Aucun GPU visible alors que DDP/gpu est demandé. Vérifie l'allocation SLURM.") + + trainer = Trainer( + accelerator="gpu", + devices=devices, # 1 GPU par process (PL mappe LOCAL_RANK -> CUDA) + num_nodes=num_nodes, + strategy=DDPStrategy( + process_group_backend="nccl", # explicite + find_unused_parameters=False, # évite des allreduces non appariés + static_graph=True, # exige même graphe/chemin à chaque itération + timeout=timedelta(minutes=15), # fail > deadlock + ), + precision="bf16-mixed", # <-- AJOUT: H100 supporte BF16 nativement + max_epochs=num_epochs, + num_sanity_val_steps=0, + callbacks=[ + ColabProgressBar(), + checkpoint_callback, + lr_monitor, + ], + log_every_n_steps=1, + gradient_clip_val=1.0, + # --- SMOKE TEST: borner provisoirement le # de batches --- + #limit_train_batches=5, # <--- TEMP pour diagnostiquer (remets 1.0 après) + #limit_val_batches=2, # <--- TEMP remettre a 2 + ) + print_dist_info("trainer_built") # Resuming training from a checkpoint if model_path and os.path.isfile(model_path) and resume_training: @@ -266,6 +361,7 @@ def main(argv): val_dataloaders=valid_dataloader, ) + print_dist_info("fit_done") if __name__ == "__main__": torch.multiprocessing.set_start_method("spawn") diff --git a/graphphysics/training/lightning_module.py b/graphphysics/training/lightning_module.py index 9dff966..2438364 100644 --- a/graphphysics/training/lightning_module.py +++ b/graphphysics/training/lightning_module.py @@ -76,6 +76,10 @@ def __init__( self.model = get_simulator(param=parameters, model=processor, device=device) + for m in getattr(self.model, "processor_list", []): + if hasattr(m, "use_activation_checkpointing"): + m.use_activation_checkpointing = True + self.loss, self.loss_name = get_loss(param=parameters) logger.info(f"Using loss {self.loss_name}") self.is_multiloss = False @@ -102,6 +106,9 @@ def __init__( self.last_val_prediction = None self.last_previous_data_prediction = None + self._last_val_num_nodes = None + self._last_pred_num_nodes = None + self.use_previous_data = use_previous_data self.previous_data_start = previous_data_start self.previous_data_end = previous_data_end @@ -116,10 +123,44 @@ def __init__( self.last_pred_prediction = None self.last_previous_data_pred_prediction = None + # === DIAG: hooks concis pour voir où ça bloque === + def on_fit_start(self): + if self.trainer.is_global_zero: + print("[diag] fit_start", flush=True) + + def on_train_epoch_start(self): + if self.trainer.is_global_zero: + print(f"[diag] epoch_start={self.current_epoch}", flush=True) + + def on_train_batch_start(self, batch, batch_idx): + # on log uniquement le tout 1er batch pour éviter le spam + if batch_idx == 0 and self.trainer.is_global_zero: + print("[diag] first_batch_start", flush=True) + + def on_after_backward(self): + # valider que le 1er backward/allreduce passe + if self.global_step == 0 and self.trainer.is_global_zero: + print("[diag] first_backward_done", flush=True) + + def on_train_batch_end(self, outputs, batch, batch_idx): + if batch_idx == 0 and self.trainer.is_global_zero: + print("[diag] first_batch_end", flush=True) + + def on_train_epoch_end(self): + if self.trainer.is_global_zero: + print(f"[diag] epoch_end={self.current_epoch}", flush=True) + ########################################################## + def forward(self, graph: Batch): return self.model(graph) def training_step(self, batch: Batch): + if self.global_step == 0 and self.trainer.is_global_zero: + print( + f"[diag] cuda reserved={torch.cuda.memory_reserved()/1e9:.2f} GB " + f"allocated={torch.cuda.memory_allocated()/1e9:.2f} GB", + flush=True, + ) batch = batch.to(self.device, non_blocking=True) node_type = batch.x[:, self.model.node_type_index] network_output, target_delta_normalized, _ = self.model(batch) @@ -145,9 +186,10 @@ def training_step(self, batch: Batch): on_step=True, on_epoch=True, prog_bar=False, + sync_dist=True, ) self.log( - "train_multiloss", loss, on_step=True, on_epoch=True, prog_bar=True + "train_multiloss", loss, on_step=True, on_epoch=True, prog_bar=True, sync_dist=True ) else: # Will raise an error if the single loss needs physical outputs. @@ -166,6 +208,7 @@ def training_step(self, batch: Batch): on_step=True, on_epoch=True, prog_bar=True, + sync_dist=True, ) return loss @@ -182,6 +225,18 @@ def _save_trajectory_to_xdmf( init_mesh = convert_to_meshio_vtu(trajectory[0], add_all_data=True) points = init_mesh.points cells = init_mesh.cells + + # --- Option: without time series (single frame only) --- + target_same_frame: bool = True + if getattr(self, "target_same_frame", True): + mesh = convert_to_meshio_vtu(trajectory[0], add_all_data=True) + meshio.write(xdmf_filename, mesh) + logger.info( + f"[No Time Series] Single frame saved at {xdmf_filename}" + ) + return + + try: with meshio.xdmf.TimeSeriesWriter(xdmf_filename) as writer: # Write the mesh (points and cells) once @@ -216,6 +271,12 @@ def _reset_validation_trajectory(self): def _make_prediction(self, batch, last_prediction, last_previous_data_prediction): batch = batch.clone() # Prepare the batch for the current step + N = batch.x.shape[0] + # reset history if graph size changed + if last_prediction is not None and last_prediction.shape[0] != N: + last_prediction = None + last_previous_data_prediction = None + if last_prediction is not None: # Update the batch with the last prediction batch.x[:, self.model.output_index_start : self.model.output_index_end] = ( @@ -240,7 +301,8 @@ def _make_prediction(self, batch, last_prediction, last_previous_data_prediction last_prediction = predicted_outputs if self.use_previous_data: last_previous_data_prediction = predicted_outputs - current_output - + # add predic velocity to batch + batch.x[:,0:3] = predicted_outputs return ( batch, predicted_outputs, @@ -255,7 +317,12 @@ def validation_step(self, batch: Batch, batch_idx: int): if batch.traj_index > self.current_val_trajectory: self._reset_validation_trajectory() self.step_counter = 0 - + # Also reset the carry if num_nodes changes within a trajectory + if self._last_val_num_nodes is not None and self._last_val_num_nodes != batch.x.shape[0]: + self.last_val_prediction = None + self.last_previous_data_prediction = None + + self._last_val_num_nodes = batch.x.shape[0] ( batch, predicted_outputs, @@ -270,26 +337,30 @@ def validation_step(self, batch: Batch, batch_idx: int): self.trajectory_to_save.append(batch) node_type = batch.x[:, self.model.node_type_index] - self.val_step_outputs.append(predicted_outputs.cpu()) - self.val_step_targets.append(target.cpu()) + #self.val_step_outputs.append(predicted_outputs.cpu()) + #self.val_step_targets.append(target.cpu()) + self.val_step_outputs.append(predicted_outputs.detach()) # rester sur GPU + self.val_step_targets.append(target.detach()) # rester sur GPU val_loss = self.val_loss( target, predicted_outputs, node_type, masks=self.loss_masks, ) - self.log("val_loss", val_loss, on_step=True, on_epoch=True, prog_bar=True) + self.log("val_loss", val_loss, on_step=True, on_epoch=True, prog_bar=True, sync_dist=True) # compute RMSE for the first step if self.step_counter == 0: squared_diff = (predicted_outputs - target) ** 2 - rmse = torch.sqrt(squared_diff.mean()).detach().cpu() + #rmse = torch.sqrt(squared_diff.mean()).detach().cpu() + rmse = torch.sqrt(squared_diff.mean()).detach() # rester sur GPU self.first_step_losses.append(rmse) self.step_counter += 1 def _reset_validation_epoch_end(self): self.val_step_outputs.clear() self.val_step_targets.clear() + self._last_val_num_nodes = None self.current_val_trajectory = 0 self.last_val_prediction = None self.last_previous_data_prediction = None @@ -297,7 +368,11 @@ def _reset_validation_epoch_end(self): self.step_counter = 0 self.first_step_losses = [] + ''' def on_validation_epoch_end(self): + # n’exécuter la logique que sur le rank global 0 + if getattr(self.trainer, "is_global_zero", False) is False: + return # Concatenate outputs and targets predicteds = torch.cat(self.val_step_outputs, dim=0) targets = torch.cat(self.val_step_targets, dim=0) @@ -312,13 +387,14 @@ def on_validation_epoch_end(self): on_step=False, on_epoch=True, prog_bar=True, + sync_dist=True, ) # Compute RMSE for the first step if self.first_step_losses: mean_first_step_loss = torch.stack(self.first_step_losses).mean().item() self.log( - "val_1step_rmse", mean_first_step_loss, on_epoch=True, prog_bar=True + "val_1step_rmse", mean_first_step_loss, on_epoch=True, prog_bar=True, sync_dist=True ) # Save trajectory graphs @@ -336,7 +412,67 @@ def on_validation_epoch_end(self): # Clear stored outputs self._reset_validation_epoch_end() + ''' + def on_validation_epoch_end(self): + # >>> Ne PAS sortir tôt sur les non-zero ranks si on utilise sync_dist=True <<< + # (on veut que tous les ranks exécutent les self.log(..., sync_dist=True)) + device = self.device + # 1) concat locales (par-rank) + predicteds = torch.cat(self.val_step_outputs, dim=0) if self.val_step_outputs else None + targets = torch.cat(self.val_step_targets, dim=0) if self.val_step_targets else None + + + + # 2) calc locales puis log avec sync_dist=True (Lightning fera la réduction) + if predicteds is not None and targets is not None: + # sécurité si jamais quelque chose est revenu sur CPU + if predicteds.device.type != "cuda": + predicteds = predicteds.to(device, non_blocking=True) + if targets.device.type != "cuda": + targets = targets.to(device, non_blocking=True) + + squared_diff = (predicteds - targets) ** 2 + all_rollout_rmse = torch.sqrt(squared_diff.mean()) + # ### IMPORTANT: laisser sync_dist=True mais appeler depuis TOUS les ranks + self.log( + "val_all_rollout_rmse", + all_rollout_rmse, + on_step=False, + on_epoch=True, + prog_bar=True, + sync_dist=True, + ) + if self.first_step_losses: + mean_first_step_loss = torch.stack(self.first_step_losses).mean() + if mean_first_step_loss.device.type != "cuda": + mean_first_step_loss= m.to(device, non_blocking=True) + # ### idem, log sur tous les ranks + self.log( + "val_1step_rmse", + mean_first_step_loss, + on_epoch=True, + prog_bar=True, + sync_dist=True, + ) + + # 3) Sauvegardes disque UNIQUEMENT sur rank 0 (I/O non distribuée) + if getattr(self.trainer, "is_global_zero", False): + save_dir = os.path.join("meshes", f"epoch_{self.current_epoch}") + self._save_trajectory_to_xdmf( + self.trajectory_to_save, + save_dir, + self._get_traj_savename( + self.trajectory_to_save, + self.current_val_trajectory, + prefix=f"graph_epoch_{self.current_epoch}", + ), + timestep=self.timestep, + ) + + # 4) reset buffers sur TOUS les ranks (pour éviter de trainer des tensors entre epochs) + self._reset_validation_epoch_end() + def configure_optimizers(self): """Initialize the optimizer""" opt = torch.optim.AdamW( @@ -381,6 +517,12 @@ def predict_step(self, batch: Batch): ) # reset self._reset_prediction_trajectory() + self._last_pred_num_nodes = None + + # If graph size changed inside a trajectory, drop the carry + if self._last_pred_num_nodes is not None and self._last_pred_num_nodes != batch.x.shape[0]: + self.last_pred_prediction = None + self.last_previous_data_pred_prediction = None # predict ( @@ -392,6 +534,7 @@ def predict_step(self, batch: Batch): ) = self._make_prediction( batch, self.last_pred_prediction, self.last_previous_data_pred_prediction ) + self._last_pred_num_nodes = batch.x.shape[0] self.prediction_trajectory.append(batch) def _reset_predict_epoch_end(self): diff --git a/graphphysics/training/parse_parameters.py b/graphphysics/training/parse_parameters.py index 42028fa..ad9921e 100644 --- a/graphphysics/training/parse_parameters.py +++ b/graphphysics/training/parse_parameters.py @@ -177,6 +177,7 @@ def get_dataset( khop = dataset_params.get("khop", 1) new_edges_ratio = dataset_params.get("new_edges_ratio", 0) extension = dataset_params.get("extension", "") + target_same_frame=dataset_params["target_same_frame"] world_pos_parameters = None if khop > 1: diff --git a/graphphysics/utils/meshio_mesh.py b/graphphysics/utils/meshio_mesh.py index c82557b..e9295de 100644 --- a/graphphysics/utils/meshio_mesh.py +++ b/graphphysics/utils/meshio_mesh.py @@ -44,10 +44,10 @@ def convert_to_meshio_vtu(graph: Data, add_all_data: bool = False) -> meshio.Mes .numpy() .T ) - cells = [ - ("tetra" if getattr(graph, "tetra", None) is not None else "triangle", faces) - ] - + #cells = [ + # ("tetra" if getattr(graph, "tetra", None) is not None else "triangle", faces) + #] + cells = [("line", faces)] # my case !!!!!!!! # Create Meshio mesh mesh = meshio.Mesh(vertices, cells) diff --git a/graphphysics/utils/torch_graph.py b/graphphysics/utils/torch_graph.py index 986c68f..b1d73ad 100644 --- a/graphphysics/utils/torch_graph.py +++ b/graphphysics/utils/torch_graph.py @@ -173,7 +173,18 @@ def meshdata_to_graph( ) if cells.shape[0] == 3: face = cells - + if cells.shape[0] == 2: #case where cells are edges + face = cells + # force Data with edge_index=face + return Data( + x=node_features, + edge_index=face, + face=face, + tetra=tetra, + y=target_features, + pos=torch.tensor(points, dtype=torch.float32), + id=id, + ) return Data( x=node_features, face=face, diff --git a/nose_training.json b/nose_training.json new file mode 100644 index 0000000..53efded --- /dev/null +++ b/nose_training.json @@ -0,0 +1,41 @@ +{ + "dataset": { + "extension": "xdmf", + "xdmf_folder": "dataset/train", + "meta_path": "dataset/nose.json", + "use_previous_data": false, + "target_same_frame": true, + "khop": 1 + }, + "model": { + "type": "transformer", + "message_passing_num": 10, + "hidden_size": 64, + "node_input_size": 3, + "output_size": 3, + "edge_input_size": 0, + "num_heads": 4, + "use_previous_data": false + }, + "index": { + "feature_index_start": 0, + "feature_index_end": 3, + "output_index_start": 0, + "output_index_end": 3, + "node_type_index": 3 + }, + "transformations": { + "preprocessing": { + "noise": [10.0, 10.0, 1.0, 10.0, 10.0, 1.0], + "noise_index_start": [0, 1, 2, 4, 5, 6], + "noise_index_end": [1, 2, 3, 5, 6, 7], + "masking": 0, + "use_previous_data": false + }, + "world_pos_parameters": { + "use": false, + "world_pos_index_start": 0, + "world_pos_index_end": 3 + } + } +} diff --git a/train.sh b/train.sh index 63ad06b..af30977 100644 --- a/train.sh +++ b/train.sh @@ -1,10 +1,15 @@ +echo "=== GPU INFO ===" +nvidia-smi +echo "================" + python -m graphphysics.train \ - --training_parameters_path=mock_training.json \ + --training_parameters_path=nose_training.json \ --num_epochs=1 \ --init_lr=0.001 \ --batch_size=1 \ --warmup=500 \ --num_workers=0 \ --prefetch_factor=0 \ - --model_save_name=model \ - --no_edge_feature + --model_save_name=model.ckpt \ + --no_edge_feature \ + --use_previous_data=false