MCPcopy Create free account

hub / github.com/KellerJordan/modded-nanogpt / functions

Functions2,320 in github.com/KellerJordan/modded-nanogpt

⨍ Functions 2,320◇ Types & classes 557↳ Endpoints 41
↓ 6 callersMethod__init__
(self, dim)
records/track_1_short/2024-12-04_ValueEmbed/train_gpt2.py:196
↓ 6 callersMethod__init__
(self, dim: int)
records/track_3_optimization/results/20260513_shampoo_1_4_power/train_gpt_shampoo.py:116
↓ 6 callersMethod__post_init__
(self)
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py:105
↓ 6 callersMethod_shampoo_optim_factory
( distributed_config: FullyShardDistributedConfig | SingleDeviceDistributedConfig, )
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/gpu_tests/shampoo_fully_shard_distributor_test.py:119
↓ 6 callersMethod_shampoo_optim_factory
( distributed_config: DDPDistributedConfig | FullyShardDistributedConfig | HybridShard
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/gpu_tests/shampoo_hybrid_shard_lossless_distributor_test.py:122
↓ 6 callersMethod_shampoo_optim_factory
( distributed_config: FSDPDistributedConfig | SingleDeviceDistributedConfig, )
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/gpu_tests/shampoo_fsdp_distributor_test.py:114
↓ 6 callersFunctionapply_tail_extrapolated_weights
(step: int, gamma: float | None = None)
records/track_3_optimization/results/20260520_tail_refinterp_2900/train_gpt_tail_refinterp_2900.py:1338
↓ 6 callersFunctionapply_tail_velocity_weights
(step: int, gamma: float | None = None)
records/track_3_optimization/results/20260520_tail_refinterp_2900/train_gpt_tail_refinterp_2900.py:1357
↓ 6 callersFunctioncompute_validation_loss
()
records/track_3_optimization/results/20260520_tail_refinterp_2900/train_gpt_tail_refinterp_2900.py:1577
↓ 6 callersFunctionget_dtype_size
Return the size (bytes) of a given data type.
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/utils/shampoo_utils.py:199
↓ 6 callersFunctionmatrix_eigendecomposition
Compute the eigendecomposition of a symmetric matrix. Args: A (Tensor): The input symmetric matrix. epsilon (float): Adds epsilon
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/preconditioner/matrix_functions.py:680
↓ 6 callersFunctionmatrix_inverse_root_from_eigendecomposition
Compute A^(-1/root) from eigendecomposition A = Q diag(L) Q^T. This function computes A^(-1/root) given the eigendecomposition of A = Q diag(L) Q
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/preconditioner/matrix_functions.py:172
↓ 6 callersFunctionnorm
(x: Tensor)
train_gpt.py:945
↓ 6 callersFunctionnorm
(x)
records/track_2_medium/2024-12-31_Target350M/train_gpt.py:131
↓ 6 callersFunctionnorm
(x: Tensor)
records/track_2_medium/2025-11-12_BlockMaskRedundantOp/train_gpt_medium.py:119
↓ 6 callersFunctionnorm
(x)
records/track_1_short/2024-12-04_ValueEmbed/train_gpt2.py:128
↓ 6 callersFunctionnorm
(x: Tensor)
records/track_1_short/2025-12-11_NorMuonOptimsAndFixes/profiler-example-traces/train_gpt-profiler-example.py:823
↓ 6 callersFunctionprint0
(s, console=False, log=True)
records/track_3_optimization/train_gpt_simple.py:233
↓ 6 callersFunctionprint0
(s, console=False, log=True)
records/track_3_optimization/results/20260519_dynmuon/train_gpt_simple_dynmuon.py:378
↓ 6 callersFunctionprint0
(s, console=False, log=True)
records/track_3_optimization/results/20260501_muonsq/train_gpt_simple_muonsq.py:245
↓ 6 callersFunctionprint0
(s, console=False, log=True)
records/track_3_optimization/results/20260430_muonh/train_gpt_simple_muonh.py:253
↓ 6 callersFunctionprint0
(s, console=False, log=True)
records/track_3_optimization/results/20260611_muonh_tuned_aux_3250/train_gpt_simple_muonh_tuned_aux_3250.py:253
↓ 6 callersFunctionprint0
(s, console=False, log=True)
records/track_3_optimization/results/20260513_shampoo_1_4_power/train_gpt_shampoo.py:198
↓ 6 callersFunctionprint0
(s, console=False, log=True)
records/track_3_optimization/results/20260527_psgd/train_psgd.py:280
↓ 6 callersFunctionprint0
(s, console=False, log=True)
records/track_3_optimization/results/20260505_newton_muon/train_gpt_simple_newton_muon.py:388
↓ 6 callersFunctionregister_resolvers
Register all custom Hydra resolvers. Call once at startup.
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/examples/resolvers.py:51
↓ 6 callersFunctionset_seed
(seed: int)
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/examples/utils.py:74
↓ 6 callersMethodupdate_params
Update params stored inside this distributor according to the input search directions argument. Args: blocked_search_directions (
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/shampoo_distributor.py:337
↓ 5 callersMethod__init__
(self, dim: int)
records/track_2_medium/2025-11-12_BlockMaskRedundantOp/train_gpt_medium.py:179
↓ 5 callersMethod__init__
(self, config)
records/track_1_short/2024-10-09_SOAP/train_gpt2.py:508
↓ 5 callersMethod__init__
(self, config)
records/track_1_short/2024-10-14_ModernArch/train_gpt2.py:167
↓ 5 callersMethod__init__
(self, config)
records/track_1_short/2024-10-10_Muon/train_gpt2.py:174
↓ 5 callersMethod_construct_local_block_info_list_with_params
( self, params: Iterable[Tensor], rank: int | None = None )
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/shampoo_distributor.py:369
↓ 5 callersMethod_construct_model_and_distributor
( self, )
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/gpu_tests/distributor_test_utils.py:42
↓ 5 callersMethod_init_distributed
(self)
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/gpu_tests/shampoo_ddp_distributor_test.py:73
↓ 5 callersMethod_init_distributed
(self)
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/gpu_tests/shampoo_ddp_distributor_test.py:746
↓ 5 callersMethodget_forward_args
(self)
train_gpt.py:1909
↓ 5 callersFunctionget_model_and_loss_fn
( device: torch.device, out_channels: int = 64, disable_linear_bias: bool = False, )
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/examples/utils.py:93
↓ 5 callersFunctiongram_frobenius_norm_estimate
(G: Tensor, keepdim: bool = False, eps: float = 1e-10)
records/track_3_optimization/results/20260611_tailema_2720_submission/train_gpt_tailema_2720.py:326
↓ 5 callersFunctiongram_frobenius_norm_estimate
(G: Tensor, keepdim: bool = False, eps: float = 1e-10)
records/track_3_optimization/results/20260609_soap_f1_auxb2_clean/train_gpt_clean_SOTA.py:234
↓ 5 callersFunctiongram_frobenius_norm_estimate
(G: Tensor, keepdim: bool = False, eps: float = 1e-10)
records/track_3_optimization/results/20260619_cwd_rowfloor_tailema/train_gpt_cwd_SOTA.py:262
↓ 5 callersFunctiongram_frobenius_norm_estimate
(G: Tensor, keepdim: bool = False, eps: float = 1e-10)
records/track_3_optimization/results/20260611_tailema_2730_submission/train_gpt_tailema_2730.py:325
↓ 5 callersFunctiongram_frobenius_norm_estimate
(G: Tensor, keepdim: bool = False, eps: float = 1e-10)
records/track_3_optimization/results/20260611_tailema_2730_submission/ablation/combo_2740/train_gpt_combo_2740.py:308
↓ 5 callersMethodinit_optimizer_module
(self)
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/utils/tests/optimizer_modules_test.py:29
↓ 5 callersMethodmerge_dims
Merges dimensions of the gradient tensor till the product of the dimensions is less than or equal to max_precond_dim.
records/track_1_short/2024-10-09_SOAP/train_gpt2.py:79
↓ 5 callersFunctionparse_float_list
(raw: str)
records/track_3_optimization/results/20260520_tail_refinterp_2900/train_gpt_tail_refinterp_2900.py:172
↓ 5 callersMethodreset
(self)
train_gpt.py:998
↓ 5 callersMethodreset
(self)
records/track_1_short/2025-12-11_NorMuonOptimsAndFixes/profiler-example-traces/train_gpt-profiler-example.py:854
↓ 5 callersFunctionshampoo_comm_profiler
Context manager that profiles communication operations in Shampoo distributors. Args: name (str): The name to use for profiling (e.g., "C
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/shampoo_dist_utils.py:19
↓ 5 callersFunctionshould_soap_param
(name: str)
records/track_3_optimization/results/20260520_rre_extrapolation_pr300_2925/train_gpt_simple_rre_pr300_2925.py:401
↓ 5 callersMethodstep
(self)
records/track_1_short/2025-12-11_NorMuonOptimsAndFixes/profiler-example-traces/train_gpt-profiler-example.py:555
↓ 4 callersMethod_construct_composable_block_ids
Construct composable block ids. Args: param_index (int): Index of the parameter in self._param_group[PARAMS]. block_i
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/shampoo_distributor.py:102
↓ 4 callersMethod_convert_float_to_float
(src: torch.Tensor, dest: torch.Tensor)
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/utils/shampoo_quantization.py:112
↓ 4 callersFunction_linear_ramp
(step: int, start_step: int, end_step: int)
records/track_3_optimization/results/20260520_rre_extrapolation_pr300_2925/train_gpt_simple_rre_pr300_2925.py:618
↓ 4 callersFunction_linear_ramp
(step: int, start_step: int, end_step: int)
records/track_3_optimization/results/20260529_tail_phase_readout_2850/train_gpt_tail_phase_readout_2850.py:645
↓ 4 callersFunction_macro_val_loss
()
records/track_3_optimization/results/20260529_tail_phase_readout_2850/train_gpt_tail_phase_readout_2850.py:2056
↓ 4 callersMethod_shampoo_optim_factory
( distributed_config: FullyShardDistributedConfig | SingleDeviceDistributedConfig, )
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/gpu_tests/shampoo_fully_shard_lossless_distributor_test.py:120
↓ 4 callersFunction_sparse_comms_active
()
train_gpt.py:252
↓ 4 callersMethod_split_tensor_block_recovery
Chunks flattened tensor in order to re-construct valid blocks with respect to the original multi-dimensional tensor shape and parameter bounda
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/shampoo_hsdp_distributor.py:652
↓ 4 callersMethod_train_step
Perform a single training step and return the loss (detached).
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/gpu_tests/shampoo_checkpoint_test.py:298
↓ 4 callersFunctiondistribute_buffer_sizes
Distribute given param blocks across ranks in a group. Param blocks are distributed such that the total assigned load of each rank is as even as
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/utils/shampoo_utils.py:273
↓ 4 callersFunctiondistributed_data_generator
(filename_pattern: str, num_tokens: int, max_seq_len: int, grad_accum_steps: int = 1, align_to_bos: bool = Tru
train_gpt_medium.py:1324
↓ 4 callersFunctiondistributed_data_generator
(filename_pattern: str, num_tokens: int, max_seq_len: int, grad_accum_steps: int = 1, align_to_bos: bool = Tru
train_gpt.py:1614
↓ 4 callersFunctiondistributed_data_generator
(filename_pattern: str, num_tokens: int, max_seq_len: int, grad_accum_steps: int = 1, align_to_bos: bool = Tru
records/track_1_short/2025-12-11_NorMuonOptimsAndFixes/profiler-example-traces/train_gpt-profiler-example.py:1238
↓ 4 callersMethodflush
(self)
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/examples/loss_metrics.py:96
↓ 4 callersMethodfrom_block
Creates a BaseShampooKroneckerFactorsState object for a given block. Args: block_info (BlockInfo): Information about the
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py:74
↓ 4 callersMethodget_forward_args
(self)
train_gpt_medium.py:1523
↓ 4 callersFunctionget_ws
(step: int)
train_gpt_medium.py:1409
↓ 4 callersMethodload_state_dict
Load optimizer state from a dict.
train_gpt.py:696
↓ 4 callersMethodlog_global_metrics
(self)
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/examples/loss_metrics.py:81
↓ 4 callersMethodlookup
(self, step: int)
train_gpt.py:1773
↓ 4 callersFunctionmatrix_orthogonalization
Compute the orthogonalization of a matrix. Args: A (Tensor): The input matrix. orthogonalization_config (OrthogonalizationConfig)
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/preconditioner/matrix_functions.py:893
↓ 4 callersFunctionnext_multiple_of_n
(v: float | int, *, n: int)
train_gpt.py:1141
↓ 4 callersMethodprecondition
( self, masked_grad_list: tuple[Tensor, ...] )
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/preconditioner/preconditioner_list.py:77
↓ 4 callersMethodquantize_mlp_fp8
Refresh the FP8 copy of the MLP up-projection weights after optimizer steps.
train_gpt.py:1332
↓ 4 callersMethodreset
(self)
train_gpt_medium.py:874
↓ 4 callersFunctionsoap_precondition_momentum
(update, state, beta2=SOAP_BETA2, eps=1e-8, blend=SOAP_BLEND, denom_floor_ratio
records/track_3_optimization/results/20260529_tail_phase_readout_2850/train_gpt_tail_phase_readout_2850.py:614
↓ 4 callersFunctionsoap_update_preconditioner
(grad, state, shampoo_beta=SOAP_BETA2, precondition_frequency=SOAP_PRECONDITION_FREQUENCY)
records/track_3_optimization/results/20260529_tail_phase_readout_2850/train_gpt_tail_phase_readout_2850.py:632
↓ 4 callersMethodstart
(self)
train_gpt_medium.py:1260
↓ 4 callersMethodwrap_model
( self, model: nn.Module, local_rank: int, backend: str, device_mesh:
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/examples/parallelism.py:80
↓ 3 callersMethod__init__
(self, vocab_size: int, num_layers: int, num_heads: int, head_dim: int, model_dim: int, max_seq_len: int)
train_gpt.py:1152
↓ 3 callersMethod__post_init__
(self)
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/shampoo_types.py:977
↓ 3 callersFunction_checkpoint_model_state
(path: str)
records/track_3_optimization/results/20260529_tail_phase_readout_2850/train_gpt_tail_phase_readout_2850.py:2009
↓ 3 callersMethod_construct_model
( post_model_decoration: Callable[[nn.Module], nn.Module] = lambda x: x, )
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/gpu_tests/shampoo_hybrid_shard_distributor_test.py:72
↓ 3 callersMethod_construct_model
( post_model_decoration: Callable[[nn.Module], nn.Module] = lambda x: x, distributed_config: H
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/gpu_tests/shampoo_hsdp_distributor_test.py:66
↓ 3 callersMethod_eye_
(x, value)
records/track_3_optimization/results/20260505_newton_muon/train_gpt_simple_newton_muon.py:251
↓ 3 callersMethod_get_params_or_grads
( self, get_grad: Literal[True] )
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/shampoo_distributor.py:139
↓ 3 callersMethod_get_train_and_eval_interp_coeffs
( iterate_averaging_config: IterateAveragingConfig | None, lr: torch.Tensor | None, st
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributed_shampoo.py:1126
↓ 3 callersMethod_instantiate_preconditioner_list
(self, **kwargs: Any)
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/preconditioner/tests/adagrad_preconditioner_list_test.py:37
↓ 3 callersMethod_is_adam_step
Adam params are only updated on odd steps.
train_gpt.py:1917
↓ 3 callersFunction_load_data_shard
(file: Path)
train_gpt_medium.py:1223
↓ 3 callersFunction_load_data_shard
(file: Path)
train_gpt.py:1515
↓ 3 callersFunction_load_data_shard
(file: Path)
records/track_1_short/2025-12-11_NorMuonOptimsAndFixes/profiler-example-traces/train_gpt-profiler-example.py:1137
↓ 3 callersFunction_load_data_shard
(file: Path)
records/track_3_optimization/results/20260520_tail_refinterp_2900/train_gpt_tail_refinterp_2900.py:237
↓ 3 callersFunction_macro_load_candidate
(base_state: dict, delta1: dict, delta2: dict, alpha: Tensor, groups_by_name: dict[str, str] | None = None)
records/track_3_optimization/results/20260529_tail_phase_readout_2850/train_gpt_tail_phase_readout_2850.py:2014
↓ 3 callersFunction_macro_score_calib
(calib_batches)
records/track_3_optimization/results/20260529_tail_phase_readout_2850/train_gpt_tail_phase_readout_2850.py:2046
↓ 3 callersMethod_merge_and_block_gradients
Split, merge, and block gradients. Returns: local_masked_blocked_grads (tuple[Tensor, ...]): Local gradients with grad not None.
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/shampoo_hsdp_distributor.py:518
↓ 3 callersMethod_merge_and_block_with_params
Merges small dimensions and blocks parameters into manageable chunks. This function processes each parameter tensor by: 1. Merging sm
records/track_3_optimization/results/20260513_shampoo_1_4_power/distributed_shampoo/distributor/shampoo_distributor.py:174
↓ 3 callersFunction_named_param_buffer_state
(buffer_dict: dict)
records/track_3_optimization/results/20260520_tail_refinterp_2900/train_gpt_tail_refinterp_2900.py:1458
↓ 3 callersFunction_ns_inner
(X: Tensor)
records/track_3_optimization/results/20260520_tail_refinterp_2900/train_gpt_tail_refinterp_2900.py:394
← previousnext →101–200 of 2,320, ranked by callers