| {"__qualname__": "kauldron.kd:train.Trainer", "workdir": "/users/staff/dmi-dmi/miele0000/diffuGemma-pt/xp_reweighted_ce_unsloth", "evals": {"gemma_checkpointer": {"__qualname__": "gemma.diffusion.hackable_diffusion_adapter.hd.gemma_checkpointer:GemmaCheckpointFormatter", "run": {"__qualname__": "kauldron.kd:evals.StandaloneEveryCheckpoint", "requirements": {"__qualname__": "xmanager.xm:JobRequirements"}, "executor": {"__qualname__": "xmanager.xm_abc:Borg", "scheduling": {"__qualname__": "xmanager.xm_abc:BorgScheduling"}, "autopilot_params": {"__qualname__": "xmanager.xm_abc:AutopilotParams"}, "borglet_params": {"__qualname__": "xmanager.xm_abc:BorgletParams"}}}, "write_original_params": false, "write_lora_params": false, "write_fused_lora_params": true}}, "schedules": {"learning_rate": {"__qualname__": "optax:warmup_cosine_decay_schedule", "init_value": 0.0, "peak_value": 0.00015, "end_value": 1.4999999999999999e-05, "warmup_steps": 100, "decay_steps": 4000, "__id__": 2}}, "setup": {"__qualname__": "kauldron.kd:train.Setup"}, "xm_job": {"__qualname__": "kauldron.kxm:Job", "args": {}, "env_vars": {}, "files": {}, "requirements": {"__qualname__": "xmanager.xm:JobRequirements"}, "executor": {"__qualname__": "xmanager.xm_abc:Borg", "scheduling": {"__qualname__": "xmanager.xm_abc:BorgScheduling"}, "autopilot_params": {"__qualname__": "xmanager.xm_abc:AutopilotParams"}, "borglet_params": {"__qualname__": "xmanager.xm_abc:BorgletParams"}}, "debug": {"__qualname__": "kauldron.kxm:Debug"}, "interpreter_info": {"__qualname__": "kauldron.kxm:InterpreterInfo"}}, "seed": 42, "aux": {"vocab_size": 262144, "sft_variant": "reweighted-ce", "schedule": {"__qualname__": "hackable_diffusion.hd:corruption.RFSchedule", "__id__": 0}, "corruption_process": {"__qualname__": "hackable_diffusion.hd:corruption.CategoricalProcess.uniform_process", "num_categories": 262144, "schedule": {"__qualname__": "hackable_diffusion.hd:corruption.RFSchedule", "__id__": 0}, "__id__": 1}, "prompt_len": 256, "num_canvases": 1, "canvas_size": 256, "use_lora": true, "lora_rank": 64, "peak_lr": 0.00015, "end_lr": 1.4999999999999999e-05, "checkpoint_every_n_steps": 1000, "eval_num_batches": null, "stop_gradient_from_denoiser_to_encoder": false, "encoder_loss_weight": 1.0, "decoder_loss_weight": 1.0, "sudoku_prompt": "<|turn>system Solve the following Sudoku puzzle. Empty cells are represented by 0. Output ONLY the solved puzzle immediately as a 9x9 grid of numbers separated by spaces. Do not include ####, explanations, or any other text.<turn|>\n<|turn>user {text}<turn|>\n<|turn>model\n"}, "sharding": {"__qualname__": "kauldron.kd:sharding.ShardingStrategy", "params": {"__qualname__": "kauldron.kd:sharding.FSDPSharding"}, "opt_state": {"__qualname__": "kauldron.kd:sharding.FSDPSharding"}}, "model": {"__qualname__": "gemma.diffusion.hackable_diffusion_adapter.hd.sft_model:SFTDiffusion", "x0": "batch.canvas", "prompt": "batch.prompt", "canvas_id": "batch.canvas_id", "canvas_mask": "batch.canvas_mask", "encoder_target": "batch.encoder_target", "encoder_target_mask": "batch.encoder_target_mask", "corruption_process": {"__qualname__": "hackable_diffusion.hd:corruption.CategoricalProcess.uniform_process", "num_categories": 262144, "schedule": {"__qualname__": "hackable_diffusion.hd:corruption.RFSchedule", "__id__": 0}, "__id__": 1}, "time_sampler": {"__qualname__": "hackable_diffusion.hd:training.time_sampling.UniformTimeSampler", "span": {"__qualname__": "hackable_diffusion.hd:jax_helpers.SafeSpan", "safety_epsilon": 0.0001}}, "gemma_network": {"__qualname__": "gemma.diffusion.hackable_diffusion_adapter.hd.lora:LoRA", "rank": 64, "model": {"__qualname__": "gemma.diffusion.hackable_diffusion_adapter.hd.hd_gemma_network:WrappedDiffusionGemmaNetwork", "gemma_model": {"__qualname__": "gemma.diffusion._models:DiffusionGemma_26B_A4B"}}, "target_modules": ["attn", "mlp2"], "alpha": 128.0}, "prompt_len": 256, "canvas_size": 256, "num_canvases": 1, "stop_gradient_from_denoiser_to_encoder": false}, "train_losses": {"diffusion_loss": {"__qualname__": "hackable_diffusion.kdiff.core:KauldronLossWrapper", "loss": {"__qualname__": "gemma.diffusion.hackable_diffusion_adapter.hd.sft_model:UniformReweightedCELoss", "schedule": {"__qualname__": "hackable_diffusion.hd:corruption.RFSchedule", "__id__": 0}, "vocab_size": 262144, "use_mask": true, "mask_key": "target_mask"}, "weight": 1.0}, "encoder_loss": {"__qualname__": "gemma.diffusion.hackable_diffusion_adapter.hd.sft_model:EncoderARLoss", "encoder_logits": "preds.encoder_logits", "encoder_target": "preds.encoder_target", "encoder_target_mask": "preds.encoder_target_mask", "weight": 1.0}}, "num_train_steps": 4000, "optimizer": {"__qualname__": "kauldron.kd:optim.partial_updates", "optimizer": {"__qualname__": "kauldron.kd:optim.named_chain", "clip": {"__qualname__": "optax:clip_by_global_norm", "max_norm": 1.0}, "adam": {"__qualname__": "optax:scale_by_adam", "b1": 0.95, "b2": 0.99, "eps": 1e-08}, "decay": {"__qualname__": "optax:add_decayed_weights", "weight_decay": 0.0001}, "lr": {"__qualname__": "optax:scale_by_learning_rate", "0": {"__qualname__": "optax:warmup_cosine_decay_schedule", "init_value": 0.0, "peak_value": 0.00015, "end_value": 1.4999999999999999e-05, "warmup_steps": 100, "decay_steps": 4000, "__id__": 2}}}, "mask": {"__qualname__": "kauldron.kd:optim.select", "0": "lora"}}, "checkpointer": {"__qualname__": "kauldron.kd:ckpts.Checkpointer", "fast": false, "save_interval_steps": 1000, "max_to_keep": 5}, "rng_streams": {"__qualname__": "kauldron.kd:train.RngStreams", "0": [{"__qualname__": "kauldron.kd:train.RngStream", "0": "default", "train": true, "eval": true}, {"__qualname__": "kauldron.kd:train.RngStream", "0": "sampling", "train": true, "eval": true}]}, "init_transform": {"__qualname__": "gemma.diffusion.hackable_diffusion_adapter.hd.gemma_checkpointer:GemmaDiffusionCheckpointLoader", "path": "/users/staff/dmi-dmi/miele0000/ckpts/diffusiongemma-26B-A4B-it"}, "train_ds": {"__qualname__": "gemma.diffusion.hackable_diffusion_adapter.data.sudoku.sudoku_data:make_sudoku_ds", "bagz_path": "gemma/diffusion/hackable_diffusion_adapter/data/sudoku/sudoku_train.bagz", "training": true, "batch_size": 4, "prompt_len": 256, "num_canvases": 1, "canvas_size": 256, "prompt_template": "<|turn>system Solve the following Sudoku puzzle. Empty cells are represented by 0. Output ONLY the solved puzzle immediately as a 9x9 grid of numbers separated by spaces. Do not include ####, explanations, or any other text.<turn|>\n<|turn>user {text}<turn|>\n<|turn>model\n"}, "eval_ds": {"__qualname__": "gemma.diffusion.hackable_diffusion_adapter.data.sudoku.sudoku_data:make_sudoku_ds", "bagz_path": "gemma/diffusion/hackable_diffusion_adapter/data/sudoku/sudoku_eval.bagz", "training": false, "batch_size": 4, "prompt_len": 256, "num_canvases": 1, "canvas_size": 256, "slice_stop": 256, "prompt_template": "<|turn>system Solve the following Sudoku puzzle. Empty cells are represented by 0. Output ONLY the solved puzzle immediately as a 9x9 grid of numbers separated by spaces. Do not include ####, explanations, or any other text.<turn|>\n<|turn>user {text}<turn|>\n<|turn>model\n"}, "writer": {"__qualname__": "gemma.diffusion.hackable_diffusion_adapter.safe_writer:SafeMetricWriter"}} |